Update app.py
Browse files
app.py
CHANGED
@@ -6,16 +6,13 @@ import os
|
|
6 |
|
7 |
torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
|
8 |
|
9 |
-
model = torch.hub.load('pytorch/vision:v0.10.0', 'shufflenet_v2_x1_0', pretrained=True)
|
10 |
-
model.load_state_dict(torch.load('shufflenetv2_x1-5666bf0f80.pth'))
|
11 |
-
|
12 |
# To export
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
|
17 |
# To load
|
18 |
-
|
19 |
|
20 |
|
21 |
os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")
|
|
|
6 |
|
7 |
torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
|
8 |
|
|
|
|
|
|
|
9 |
# To export
|
10 |
+
model = torch.hub.load('pytorch/vision:v0.10.0', 'shufflenet_v2_x1_0', pretrained=True).eval()
|
11 |
+
traced_graph = torch.jit.trace(model, torch.randn(1, 3, H, W))
|
12 |
+
traced_graph.save('shufflenet.pth')
|
13 |
|
14 |
# To load
|
15 |
+
model = torch.jit.load('shufflenet.pth').eval().to(device)
|
16 |
|
17 |
|
18 |
os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")
|