czkaiweb commited on
Commit
b36c717
·
1 Parent(s): 33c60f2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -11
app.py CHANGED
@@ -11,20 +11,20 @@ import torch.nn as nn
11
  #model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=False)
12
  #checkpoint = 'https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/model_weights_mobilenet_v2_valp1trainp2.pth'
13
  #model.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=False))
14
- model = models.efficientnet_b0()
15
- num_ftrs = model.classifier[1].in_features
16
- model.classifier[1] = nn.Linear(num_ftrs, 6)
17
- device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
18
  #model = model.to(device)
19
- model.load_state_dict(torch.load('model_weights_EfficientNetB0_final.pth',map_location=device))
20
 
21
 
22
- #model = models.vgg16()
23
- #num_ftrs = model.classifier[6].in_features
24
- #model.classifier[6] = nn.Linear(num_ftrs, 6)
25
- #device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
26
- #model = model.to(device)
27
- #model.load_state_dict(torch.load('VGG16_weights_May28.pth',map_location=device))
28
  #torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
29
 
30
 
 
11
  #model = torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=False)
12
  #checkpoint = 'https://github.com/liuxiaoyuyuyu/vanGogh-and-Other-Artist/blob/main/model_weights_mobilenet_v2_valp1trainp2.pth'
13
  #model.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=False))
14
+ #model = models.efficientnet_b0()
15
+ #num_ftrs = model.classifier[1].in_features
16
+ #model.classifier[1] = nn.Linear(num_ftrs, 6)
17
+ #device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
18
  #model = model.to(device)
19
+ #model.load_state_dict(torch.load('model_weights_EfficientNetB0_final.pth',map_location=device))
20
 
21
 
22
+ model = models.vgg16()
23
+ num_ftrs = model.classifier[6].in_features
24
+ model.classifier[6] = nn.Linear(num_ftrs, 6)
25
+ device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
26
+ model = model.to(device)
27
+ model.load_state_dict(torch.load('VGG16_weights_May28.pth',map_location=device))
28
  #torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
29
 
30