Merge pull request #91 from mickaelseznec/master

Run on CPU + minor updates

Author: ClementPinard

main.py

@@ -163,14 +163,17 @@ def main():
         network_data = None
         print("=> creating model '{}'".format(args.arch))
 
-    model = models.__dict__[args.arch](network_data).cuda()
-    model = torch.nn.DataParallel(model).cuda()
-    cudnn.benchmark = True
+    model = models.__dict__[args.arch](network_data).to(device)
 
     assert(args.solver in ['adam', 'sgd'])
     print('=> setting {} solver'.format(args.solver))
-    param_groups = [{'params': model.module.bias_parameters(), 'weight_decay': args.bias_decay},
-                    {'params': model.module.weight_parameters(), 'weight_decay': args.weight_decay}]
+    param_groups = [{'params': model.bias_parameters(), 'weight_decay': args.bias_decay},
+                    {'params': model.weight_parameters(), 'weight_decay': args.weight_decay}]
+
+    if device.type == "cuda":
+        model = torch.nn.DataParallel(model).cuda()
+        cudnn.benchmark = True
+
     if args.solver == 'adam':
         optimizer = torch.optim.Adam(param_groups, args.lr,
                                      betas=(args.momentum, args.beta))

requirements.txt

@@ -5,4 +5,6 @@ spatial-correlation-sampler>=0.2.1
 tensorboardX>=1.4
 imageio
 argparse
-path.py
+path
+tqdm
+scipy

run_inference.py

@@ -71,7 +71,7 @@ def main():
     for ext in args.img_exts:
         test_files = data_dir.files('*1.{}'.format(ext))
         for file in test_files:
-            img_pair = file.parent / (file.namebase[:-1] + '2.{}'.format(ext))
+            img_pair = file.parent / (file.stem[:-1] + '2.{}'.format(ext))
             if img_pair.isfile():
                 img_pairs.append([file, img_pair])
 
@@ -103,7 +103,7 @@ def main():
         if args.upsampling is not None:
             output = F.interpolate(output, size=img1.size()[-2:], mode=args.upsampling, align_corners=False)
         for suffix, flow_output in zip(['flow', 'inv_flow'], output):
-            filename = save_path/'{}{}'.format(img1_file.namebase[:-1], suffix)
+            filename = save_path/'{}{}'.format(img1_file.stem[:-1], suffix)
             if args.output_value in['vis', 'both']:
                 rgb_flow = flow2rgb(args.div_flow * flow_output, max_value=args.max_flow)
                 to_save = (rgb_flow * 255).astype(np.uint8).transpose(1,2,0)