upgrade to pytorch 0.4.1 syntaxe

Author: Clément Pinard

main.py

@@ -4,6 +4,7 @@
 import time
 
 import torch
+import torch.nn.functional as F
 import torch.nn.parallel
 import torch.backends.cudnn as cudnn
 import torch.optim
@@ -19,7 +20,6 @@
 
 model_names = sorted(name for name in models.__dict__
                      if name.islower() and not name.startswith("__"))
-
 dataset_names = sorted(name for name in datasets.__all__)
 
 
@@ -82,6 +82,7 @@
 
 best_EPE = -1
 n_iter = 0
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 def main():
@@ -193,7 +194,8 @@ def main():
 
         # evaluate on validation set
 
-        EPE = validate(val_loader, model, epoch, output_writers)
+        with torch.no_grad():
+            EPE = validate(val_loader, model, epoch, output_writers)
         test_writer.add_scalar('mean EPE', EPE, epoch)
 
         if best_EPE < 0:
@@ -227,25 +229,23 @@ def train(train_loader, model, optimizer, epoch, train_writer):
     for i, (input, target) in enumerate(train_loader):
         # measure data loading time
         data_time.update(time.time() - end)
-        target = target.cuda(async=True)
-        input = [j.cuda() for j in input]
-        input_var = torch.autograd.Variable(torch.cat(input,1))
-        target_var = torch.autograd.Variable(target)
+        target = target.to(device)
+        input = torch.cat(input,1).to(device)
 
         # compute output
-        output = model(input_var)
+        output = model(input)
         if args.sparse:
             # Since Target pooling is not very precise when sparse,
             # take the highest resolution prediction and upsample it instead of downsampling target
-            h, w = target_var.size()[-2:]
-            output = [torch.nn.functional.upsample(output[0], (h,w)), *output[1:]]
+            h, w = target.size()[-2:]
+            output = [F.interpolate(output[0], (h,w)), *output[1:]]
 
-        loss = multiscaleEPE(output, target_var, weights=args.multiscale_weights, sparse=args.sparse)
-        flow2_EPE = args.div_flow * realEPE(output[0], target_var, sparse=args.sparse)
+        loss = multiscaleEPE(output, target, weights=args.multiscale_weights, sparse=args.sparse)
+        flow2_EPE = args.div_flow * realEPE(output[0], target, sparse=args.sparse)
         # record loss and EPE
-        losses.update(loss.data[0], target.size(0))
-        train_writer.add_scalar('train_loss', loss.data[0], n_iter)
-        flow2_EPEs.update(flow2_EPE.data[0], target.size(0))
+        losses.update(loss.item(), target.size(0))
+        train_writer.add_scalar('train_loss', loss.item(), n_iter)
+        flow2_EPEs.update(flow2_EPE.item(), target.size(0))
 
         # compute gradient and do optimization step
         optimizer.zero_grad()
@@ -278,26 +278,26 @@ def validate(val_loader, model, epoch, output_writers):
 
     end = time.time()
     for i, (input, target) in enumerate(val_loader):
-        target = target.cuda(async=True)
-        input_var = torch.autograd.Variable(torch.cat(input,1).cuda(), volatile=True)
-        target_var = torch.autograd.Variable(target, volatile=True)
+        target = target.to(device)
+        input = torch.cat(input,1).to(device)
 
         # compute output
-        output = model(input_var)
-        flow2_EPE = args.div_flow*realEPE(output, target_var, sparse=args.sparse)
+        output = model(input)
+        flow2_EPE = args.div_flow*realEPE(output, target, sparse=args.sparse)
         # record EPE
-        flow2_EPEs.update(flow2_EPE.data[0], target.size(0))
+        flow2_EPEs.update(flow2_EPE.item(), target.size(0))
 
         # measure elapsed time
         batch_time.update(time.time() - end)
         end = time.time()
 
         if i < len(output_writers):  # log first output of first batches
             if epoch == 0:
-                output_writers[i].add_image('GroundTruth', flow2rgb(args.div_flow * target[0].cpu().numpy(), max_value=10), 0)
-                output_writers[i].add_image('Inputs', input[0][0].numpy().transpose(1, 2, 0) + np.array([0.411,0.432,0.45]), 0)
-                output_writers[i].add_image('Inputs', input[1][0].numpy().transpose(1, 2, 0) + np.array([0.411,0.432,0.45]), 1)
-            output_writers[i].add_image('FlowNet Outputs', flow2rgb(args.div_flow * output.data[0].cpu().numpy(), max_value=10), epoch)
+                mean_values = torch.tensor([0.411,0.432,0.45], dtype=input.dtype).view(3,1,1)
+                output_writers[i].add_image('GroundTruth', flow2rgb(args.div_flow * target[0], max_value=10), 0)
+                output_writers[i].add_image('Inputs', (input[0,:3].cpu() + mean_values).clamp(0,1), 0)
+                output_writers[i].add_image('Inputs', (input[0,3:].cpu() + mean_values).clamp(0,1), 1)
+            output_writers[i].add_image('FlowNet Outputs', flow2rgb(args.div_flow * output[0], max_value=10), epoch)
 
         if i % args.print_freq == 0:
             print('Test: [{0}/{1}]\t Time {2}\t EPE {3}'
@@ -337,17 +337,17 @@ def __repr__(self):
 
 
 def flow2rgb(flow_map, max_value):
-    global args
-    _, h, w = flow_map.shape
-    flow_map[:,(flow_map[0] == 0) & (flow_map[1] == 0)] = float('nan')
-    rgb_map = np.ones((h,w,3)).astype(np.float32)
+    flow_map_np = flow_map.detach().cpu().numpy()
+    _, h, w = flow_map_np.shape
+    flow_map_np[:,(flow_map_np[0] == 0) & (flow_map_np[1] == 0)] = float('nan')
+    rgb_map = np.ones((3,h,w)).astype(np.float32)
     if max_value is not None:
-        normalized_flow_map = flow_map / max_value
+        normalized_flow_map = flow_map_np / max_value
     else:
-        normalized_flow_map = flow_map / (np.abs(flow_map).max())
-    rgb_map[:,:,0] += normalized_flow_map[0]
-    rgb_map[:,:,1] -= 0.5*(normalized_flow_map[0] + normalized_flow_map[1])
-    rgb_map[:,:,2] += normalized_flow_map[1]
+        normalized_flow_map = flow_map_np / (np.abs(flow_map_np).max())
+    rgb_map[0] += normalized_flow_map[0]
+    rgb_map[1] -= 0.5*(normalized_flow_map[0] + normalized_flow_map[1])
+    rgb_map[2] += normalized_flow_map[1]
     return rgb_map.clip(0,1)
 
 

models/FlowNetS.py

@@ -75,12 +75,12 @@ def __init__(self,batchNorm=True):
 
         for m in self.modules():
             if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
-                kaiming_normal(m.weight.data)
+                kaiming_normal_(m.weight, 0.1)
                 if m.bias is not None:
-                    m.bias.data.zero_()
+                    constant_(m.bias, 0)
             elif isinstance(m, nn.BatchNorm2d):
-                m.weight.data.fill_(1)
-                m.bias.data.zero_()
+                constant_(m.weight, 1)
+                constant_(m.bias, 0)
 
     def forward(self, x):
         out_conv2 = self.conv2(self.conv1(x))

multiscaleloss.py

@@ -1,5 +1,5 @@
 import torch
-import torch.nn as nn
+import torch.nn.functional as F
 
 
 def EPE(input_flow, target_flow, sparse=False, mean=True):
@@ -9,17 +9,25 @@ def EPE(input_flow, target_flow, sparse=False, mean=True):
         # invalid flow is defined with both flow coordinates to be exactly 0
         mask = (target_flow[:,0] == 0) & (target_flow[:,1] == 0)
 
-        EPE_map = EPE_map[~mask.data]
+        EPE_map = EPE_map[~mask]
     if mean:
         return EPE_map.mean()
     else:
         return EPE_map.sum()/batch_size
 
 
 def sparse_max_pool(input, size):
+    '''Downsample the input by considering 0 values as invalid.
+
+    Unfortunately, no generic interpolation mode can resize a sparse map correctly,
+    the strategy here is to use max pooling for positive values and "min pooling"
+    for negative values, the two results are then summed.
+    This technique allows sparsity to be minized, contrary to nearest interpolation,
+    which could potentially lose information for isolated data points.'''
+
     positive = (input > 0).float()
     negative = (input < 0).float()
-    output = nn.functional.adaptive_max_pool2d(input * positive, size) - nn.functional.adaptive_max_pool2d(-input * negative, size)
+    output = F.adaptive_max_pool2d(input * positive, size) - F.adaptive_max_pool2d(-input * negative, size)
     return output
 
 
@@ -31,7 +39,7 @@ def one_scale(output, target, sparse):
         if sparse:
             target_scaled = sparse_max_pool(target, (h, w))
         else:
-            target_scaled = nn.functional.adaptive_avg_pool2d(target, (h, w))
+            target_scaled = F.interpolate(target, (h, w), mode='area')
         return EPE(output, target_scaled, sparse, mean=False)
 
     if type(network_output) not in [tuple, list]:
@@ -48,5 +56,5 @@ def one_scale(output, target, sparse):
 
 def realEPE(output, target, sparse=False):
     b, _, h, w = target.size()
-    upsampled_output = nn.functional.upsample(output, size=(h,w), mode='bilinear')
+    upsampled_output = F.interpolate(output, (h,w), mode='bilinear', align_corners=False)
     return EPE(upsampled_output, target, sparse, mean=True)