added upsampling module

core/raft.py

@@ -3,11 +3,23 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from modules.update import BasicUpdateBlock, SmallUpdateBlock
-from modules.extractor import BasicEncoder, SmallEncoder
-from modules.corr import CorrBlock
+from update import BasicUpdateBlock, SmallUpdateBlock
+from extractor import BasicEncoder, SmallEncoder
+from corr import CorrBlock
 from utils.utils import bilinear_sampler, coords_grid, upflow8
 
+try:
+    autocast = torch.cuda.amp.autocast
+except:
+    # dummy autocast for PyTorch < 1.6
+    class autocast:
+        def __init__(self, enabled):
+            pass
+        def __enter__(self):
+            pass
+        def __exit__(self, *args):
+            pass
+
 
 class RAFT(nn.Module):
     def __init__(self, args):
@@ -26,7 +38,7 @@ class RAFT(nn.Module):
             args.corr_levels = 4
             args.corr_radius = 4
 
-        if not hasattr(args, 'dropout'):
+        if 'dropout' not in args._get_kwargs():
             args.dropout = 0
 
         # feature network, context network, and update block
@@ -40,6 +52,7 @@ class RAFT(nn.Module):
             self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='batch', dropout=args.dropout)
             self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim)
 
+
     def freeze_bn(self):
         for m in self.modules():
             if isinstance(m, nn.BatchNorm2d):
@@ -54,46 +67,73 @@ class RAFT(nn.Module):
         # optical flow computed as difference: flow = coords1 - coords0
         return coords0, coords1
 
-    def forward(self, image1, image2, iters=12, flow_init=None, upsample=True):
+    def upsample_flow(self, flow, mask):
+        """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """
+        N, _, H, W = flow.shape
+        mask = mask.view(N, 1, 9, 8, 8, H, W)
+        mask = torch.softmax(mask, dim=2)
+
+        up_flow = F.unfold(8 * flow, [3,3], padding=1)
+        up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)
+
+        up_flow = torch.sum(mask * up_flow, dim=2)
+        up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
+        return up_flow.reshape(N, 2, 8*H, 8*W)
+
+
+    def forward(self, image1, image2, iters=12, flow_init=None, upsample=True, test_mode=False):
         """ Estimate optical flow between pair of frames """
 
         image1 = 2 * (image1 / 255.0) - 1.0
         image2 = 2 * (image2 / 255.0) - 1.0
 
+        image1 = image1.contiguous()
+        image2 = image2.contiguous()
+
         hdim = self.hidden_dim
         cdim = self.context_dim
 
         # run the feature network
-        fmap1, fmap2 = self.fnet([image1, image2])
+        with autocast(enabled=self.args.mixed_precision):
+            fmap1, fmap2 = self.fnet([image1, image2])        
+        
+        fmap1 = fmap1.float()
+        fmap2 = fmap2.float()
         corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
 
         # run the context network
-        cnet = self.cnet(image1)
-        net, inp = torch.split(cnet, [hdim, cdim], dim=1)
-        net, inp = torch.tanh(net), torch.relu(inp)
+        with autocast(enabled=self.args.mixed_precision):
+            cnet = self.cnet(image1)
+            net, inp = torch.split(cnet, [hdim, cdim], dim=1)
+            net = torch.tanh(net)
+            inp = torch.relu(inp)
 
-        # if dropout is being used reset mask
-        self.update_block.reset_mask(net, inp)
         coords0, coords1 = self.initialize_flow(image1)
 
+        if flow_init is not None:
+            coords1 = coords1 + flow_init
+
         flow_predictions = []
         for itr in range(iters):
             coords1 = coords1.detach()
             corr = corr_fn(coords1) # index correlation volume
 
             flow = coords1 - coords0
-            net, delta_flow = self.update_block(net, inp, corr, flow)
+            with autocast(enabled=self.args.mixed_precision):
+                net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)
 
             # F(t+1) = F(t) + \Delta(t)
             coords1 = coords1 + delta_flow
-            
-            if upsample:
+
+            # upsample predictions
+            if up_mask is None:
                 flow_up = upflow8(coords1 - coords0)
-                flow_predictions.append(flow_up)
-            
             else:
-                flow_predictions.append(coords1 - coords0)
+                flow_up = self.upsample_flow(coords1 - coords0, up_mask)
+            
+            flow_predictions.append(flow_up)
 
+        if test_mode:
+            return coords1 - coords0, flow_up
+            
         return flow_predictions
-
-