add FlowNetC

Author: Clément Pinard

README.md

@@ -12,8 +12,8 @@ Two neural network models are currently provided :
 
  - **FlowNetS**
  - **FlowNetSBN**
-
-There is currently no implementation of FlowNetC as a specific Correlation layer module would need to be written (feel free to contribute !)
+ - **FlowNetC**
+ - **FlowNetCBN**
 
 ## Pretrained Models
 Thanks to [Kaixhin](https://github.com/Kaixhin) you can download a pretrained version of FlowNetS (from caffe, not from pytorch) [here](https://drive.google.com/open?id=0B5EC7HMbyk3CbjFPb0RuODI3NmM). This folder also contains trained networks from scratch.
@@ -22,13 +22,13 @@ Thanks to [Kaixhin](https://github.com/Kaixhin) you can download a pretrained ve
 Directly feed the downloaded Network to the script, you don't need to uncompress it even if your desktop environment tells you so.
 
 ### Note on networks from caffe
-These networks expect a BGR input in range [-0.5,0.5] (compared to RGB in pytorch). However, BGR order is not very important.
+These networks expect a BGR input in range `[-0.5,0.5]` (compared to RGB in pytorch). However, BGR order is not very important.
 
 ## Prerequisite
 
-	pytorch >= 0.3
+	pytorch >= 0.4.1
 	tensorboard-pytorch
-	tensorboardX
+	tensorboardX >= 1.4
 	scipy
 	argparse
 

models/FlowNetC.py

@@ -0,0 +1,135 @@
+import torch
+import torch.nn as nn
+from torch.nn.init import kaiming_normal_, constant_
+from .util import conv, predict_flow, deconv, crop_like, correlate
+
+__all__ = [
+    'flownetc', 'flownetc_bn'
+]
+
+
+class FlowNetC(nn.Module):
+    expansion = 1
+
+    def __init__(self,batchNorm=True):
+        super(FlowNetC,self).__init__()
+
+        self.batchNorm = batchNorm
+        self.conv1      = conv(self.batchNorm,   3,   64, kernel_size=7, stride=2)
+        self.conv2      = conv(self.batchNorm,  64,  128, kernel_size=5, stride=2)
+        self.conv3      = conv(self.batchNorm, 128,  256, kernel_size=5, stride=2)
+        self.conv_redir = conv(self.batchNorm, 256,   32, kernel_size=1, stride=1)
+
+        self.conv3_1 = conv(self.batchNorm, 473,  256)
+        self.conv4   = conv(self.batchNorm, 256,  512, stride=2)
+        self.conv4_1 = conv(self.batchNorm, 512,  512)
+        self.conv5   = conv(self.batchNorm, 512,  512, stride=2)
+        self.conv5_1 = conv(self.batchNorm, 512,  512)
+        self.conv6   = conv(self.batchNorm, 512, 1024, stride=2)
+        self.conv6_1 = conv(self.batchNorm,1024, 1024)
+
+        self.deconv5 = deconv(1024,512)
+        self.deconv4 = deconv(1026,256)
+        self.deconv3 = deconv(770,128)
+        self.deconv2 = deconv(386,64)
+
+        self.predict_flow6 = predict_flow(1024)
+        self.predict_flow5 = predict_flow(1026)
+        self.predict_flow4 = predict_flow(770)
+        self.predict_flow3 = predict_flow(386)
+        self.predict_flow2 = predict_flow(194)
+
+        self.upsampled_flow6_to_5 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False)
+        self.upsampled_flow5_to_4 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False)
+        self.upsampled_flow4_to_3 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False)
+        self.upsampled_flow3_to_2 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=False)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+                kaiming_normal_(m.weight, 0.1)
+                if m.bias is not None:
+                    constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                constant_(m.weight, 1)
+                constant_(m.bias, 0)
+
+    def forward(self, x):
+        x1 = x[:,:3]
+        x2 = x[:,3:]
+
+        out_conv1a = self.conv1(x1)
+        out_conv2a = self.conv2(out_conv1a)
+        out_conv3a = self.conv3(out_conv2a)
+
+        out_conv1b = self.conv1(x2)
+        out_conv2b = self.conv2(out_conv1b)
+        out_conv3b = self.conv3(out_conv2b)
+
+        out_conv_redir = self.conv_redir(out_conv3a)
+        out_correlation = correlate(out_conv3a,out_conv3b)
+
+        in_conv3_1 = torch.cat([out_conv_redir, out_correlation], dim=1)
+
+        out_conv3 = self.conv3_1(in_conv3_1)
+        out_conv4 = self.conv4_1(self.conv4(out_conv3))
+        out_conv5 = self.conv5_1(self.conv5(out_conv4))
+        out_conv6 = self.conv6_1(self.conv6(out_conv5))
+
+        flow6       = self.predict_flow6(out_conv6)
+        flow6_up    = crop_like(self.upsampled_flow6_to_5(flow6), out_conv5)
+        out_deconv5 = crop_like(self.deconv5(out_conv6), out_conv5)
+
+        concat5 = torch.cat((out_conv5,out_deconv5,flow6_up),1)
+        flow5       = self.predict_flow5(concat5)
+        flow5_up    = crop_like(self.upsampled_flow5_to_4(flow5), out_conv4)
+        out_deconv4 = crop_like(self.deconv4(concat5), out_conv4)
+
+        concat4 = torch.cat((out_conv4,out_deconv4,flow5_up),1)
+        flow4       = self.predict_flow4(concat4)
+        flow4_up    = crop_like(self.upsampled_flow4_to_3(flow4), out_conv3)
+        out_deconv3 = crop_like(self.deconv3(concat4), out_conv3)
+
+        concat3 = torch.cat((out_conv3,out_deconv3,flow4_up),1)
+        flow3       = self.predict_flow3(concat3)
+        flow3_up    = crop_like(self.upsampled_flow3_to_2(flow3), out_conv2a)
+        out_deconv2 = crop_like(self.deconv2(concat3), out_conv2a)
+
+        concat2 = torch.cat((out_conv2a,out_deconv2,flow3_up),1)
+        flow2 = self.predict_flow2(concat2)
+
+        if self.training:
+            return flow2,flow3,flow4,flow5,flow6
+        else:
+            return flow2
+
+    def weight_parameters(self):
+        return [param for name, param in self.named_parameters() if 'weight' in name]
+
+    def bias_parameters(self):
+        return [param for name, param in self.named_parameters() if 'bias' in name]
+
+
+def flownetc(data=None):
+    """FlowNetS model architecture from the
+    "Learning Optical Flow with Convolutional Networks" paper (https://arxiv.org/abs/1504.06852)
+
+    Args:
+        data : pretrained weights of the network. will create a new one if not set
+    """
+    model = FlowNetC(batchNorm=False)
+    if data is not None:
+        model.load_state_dict(data['state_dict'])
+    return model
+
+
+def flownetc_bn(data=None):
+    """FlowNetS model architecture from the
+    "Learning Optical Flow with Convolutional Networks" paper (https://arxiv.org/abs/1504.06852)
+
+    Args:
+        data : pretrained weights of the network. will create a new one if not set
+    """
+    model = FlowNetC(batchNorm=True)
+    if data is not None:
+        model.load_state_dict(data['state_dict'])
+    return model

models/FlowNetS.py

@@ -1,44 +1,13 @@
 import torch
 import torch.nn as nn
-from torch.nn.init import kaiming_normal
+from torch.nn.init import kaiming_normal_, constant_
+from .util import conv, predict_flow, deconv, crop_like
 
 __all__ = [
     'flownets', 'flownets_bn'
 ]
 
 
-def conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1):
-    if batchNorm:
-        return nn.Sequential(
-            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=False),
-            nn.BatchNorm2d(out_planes),
-            nn.LeakyReLU(0.1,inplace=True)
-        )
-    else:
-        return nn.Sequential(
-            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True),
-            nn.LeakyReLU(0.1,inplace=True)
-        )
-
-
-def predict_flow(in_planes):
-    return nn.Conv2d(in_planes,2,kernel_size=3,stride=1,padding=1,bias=False)
-
-
-def deconv(in_planes, out_planes):
-    return nn.Sequential(
-        nn.ConvTranspose2d(in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=False),
-        nn.LeakyReLU(0.1,inplace=True)
-    )
-
-
-def crop_like(input, target):
-    if input.size()[2:] == target.size()[2:]:
-        return input
-    else:
-        return input[:, :, :target.size(2), :target.size(3)]
-
-
 class FlowNetS(nn.Module):
     expansion = 1
 

models/__init__.py

@@ -1 +1,2 @@
-from .FlowNetS import *
+from .FlowNetS import *
+from .FlowNetC import *

models/util.py

@@ -0,0 +1,58 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from spatial_correlation_sampler import spatial_correlation_sample
+except ImportError as e:
+    import warnings
+    with warnings.catch_warnings():
+        warnings.filterwarnings("default", category=ImportWarning)
+        warnings.warn("failed to load custom correlation module"
+                      "which is needed for FlowNetC", ImportWarning)
+
+
+def conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1):
+    if batchNorm:
+        return nn.Sequential(
+            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=False),
+            nn.BatchNorm2d(out_planes),
+            nn.LeakyReLU(0.1,inplace=True)
+        )
+    else:
+        return nn.Sequential(
+            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True),
+            nn.LeakyReLU(0.1,inplace=True)
+        )
+
+
+def predict_flow(in_planes):
+    return nn.Conv2d(in_planes,2,kernel_size=3,stride=1,padding=1,bias=False)
+
+
+def deconv(in_planes, out_planes):
+    return nn.Sequential(
+        nn.ConvTranspose2d(in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=False),
+        nn.LeakyReLU(0.1,inplace=True)
+    )
+
+
+def correlate(input1, input2):
+    out_corr = spatial_correlation_sample(input1,
+                                          input2,
+                                          kernel_size=1,
+                                          patch_size=21,
+                                          stride=1,
+                                          padding=0,
+                                          dilation_patch=2)
+    # collate dimensions 1 and 2 in order to be treated as a
+    # regular 4D tensor
+    b, ph, pw, h, w = out_corr.size()
+    out_corr = out_corr.view(b, ph * pw, h, w)/input1.size(1)
+    return F.leaky_relu_(out_corr, 0.1)
+
+
+def crop_like(input, target):
+    if input.size()[2:] == target.size()[2:]:
+        return input
+    else:
+        return input[:, :, :target.size(2), :target.size(3)]

requirements.txt

@@ -0,0 +1,7 @@
+torch>=0.4.1
+torchvision
+numpy
+spatial-correlation-sampler
+tensorboardX>=1.4
+imageio
+argparse