Merge pull request #4 from BorisTestov/cpu-key

Cpu key

demo.py

@@ -20,6 +20,10 @@
                     help='path to the input picture')
 parser.add_argument('--outdir', default='./test_result', type=str,
                     help='path to save the predict result')
+
+parser.add_argument('--cpu', dest='cpu', action='store_true')
+parser.set_defaults(cpu=False)
+
 args = parser.parse_args()
 
 
@@ -36,10 +40,8 @@ def demo():
     ])
     image = Image.open(args.input_pic).convert('RGB')
     image = transform(image).unsqueeze(0).to(device)
-
-    model = get_fast_scnn(args.dataset, pretrained=True, root=args.weights_folder).to(device)
+    model = get_fast_scnn(args.dataset, pretrained=True, root=args.weights_folder, map_cpu=args.cpu).to(device)
     print('Finished loading model!')
-
     model.eval()
     with torch.no_grad():
         outputs = model(image)

models/fast_scnn.py

@@ -233,7 +233,7 @@ def forward(self, x):
         return x
 
 
-def get_fast_scnn(dataset='citys', pretrained=False, root='./weights', **kwargs):
+def get_fast_scnn(dataset='citys', pretrained=False, root='./weights', map_cpu=False, **kwargs):
     acronyms = {
         'pascal_voc': 'voc',
         'pascal_aug': 'voc',
@@ -244,7 +244,10 @@ def get_fast_scnn(dataset='citys', pretrained=False, root='./weights', **kwargs)
     from data_loader import datasets
     model = FastSCNN(datasets[dataset].NUM_CLASS, **kwargs)
     if pretrained:
-        model.load_state_dict(torch.load(os.path.join(root, 'fast_scnn_%s_best_model.pth' % acronyms[dataset])))
+        if(map_cpu):
+            model.load_state_dict(torch.load(os.path.join(root, 'fast_scnn_%s.pth' % acronyms[dataset]), map_location='cpu'))
+        else:
+            model.load_state_dict(torch.load(os.path.join(root, 'fast_scnn_%s.pth' % acronyms[dataset])))
     return model