t

Author: Study-is-happy

NeuFlow/backbone_v6.py

@@ -54,8 +54,7 @@ def init_pos(self, batch_size, height, width, device, amp):
         return pos[None].repeat(batch_size,1,1,1)
 
     def init_bhwd(self, batch_size, height, width, device, amp):
-        self.pos_s8 = self.init_pos(batch_size, height//8, width//8, device, amp)
-        self.pos_s16 = self.init_pos(batch_size, height//16, width//16, device, amp)
+        self.pos_s16 = self.init_pos(batch_size, height, width, device, amp)
 
     def forward(self, img):
 
@@ -75,6 +74,5 @@ def forward(self, img):
         x_16 = self.block_cat_16(torch.cat([x_16, x_16_2], dim=1))
 
         x_16 = torch.cat([x_16, self.pos_s16], dim=1)
-        x_8 = torch.cat([x_8, self.pos_s8], dim=1)
 
         return x_16, x_8

NeuFlow/backbone_v7.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn.functional as F
+
+
+class ConvBlock(torch.nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, stride, padding):
+        super(ConvBlock, self).__init__()
+
+        self.conv1 = torch.nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode='zeros', bias=False)
+
+        self.conv2 = torch.nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False)
+
+        self.relu = torch.nn.LeakyReLU(negative_slope=0.1, inplace=False)
+
+        self.norm1 = torch.nn.BatchNorm2d(out_planes)
+
+        self.norm2 = torch.nn.BatchNorm2d(out_planes)
+
+        # self.dropout = torch.nn.Dropout(p=0.1)
+
+    def forward(self, x):
+
+        # x = self.dropout(x)
+
+        x = self.relu(self.norm1(self.conv1(x)))
+        x = self.relu(self.norm2(self.conv2(x)))
+        # x = self.relu(self.conv1(x))
+        # x = self.relu(self.conv2(x))
+
+        return x
+
+class CNNEncoder(torch.nn.Module):
+    def __init__(self, feature_dim_s16, context_dim_s16, feature_dim_s8, context_dim_s8):
+        super(CNNEncoder, self).__init__()
+
+        self.block_8_1 = ConvBlock(3, feature_dim_s8 * 2, kernel_size=8, stride=4, padding=2)
+
+        self.block_8_2 = ConvBlock(3, feature_dim_s8, kernel_size=6, stride=2, padding=2)
+
+        self.block_cat_8 = ConvBlock(feature_dim_s8 * 3, feature_dim_s8 + context_dim_s8, kernel_size=3, stride=1, padding=1)
+
+        self.block_16_1 = ConvBlock(3, feature_dim_s16, kernel_size=6, stride=2, padding=2)
+
+        self.block_8_16 = ConvBlock(feature_dim_s8 + context_dim_s8, feature_dim_s16, kernel_size=6, stride=2, padding=2)
+
+        self.block_cat_16 = ConvBlock(feature_dim_s16 * 2, feature_dim_s16 + context_dim_s16 - 2, kernel_size=3, stride=1, padding=1)
+
+    def init_pos(self, batch_size, height, width, device, amp):
+        ys, xs = torch.meshgrid(torch.arange(height, dtype=torch.half if amp else torch.float, device=device), torch.arange(width, dtype=torch.half if amp else torch.float, device=device), indexing='ij')
+        ys = (ys-height/2)
+        xs = (xs-width/2)
+        pos = torch.stack([ys, xs])
+        return pos[None].repeat(batch_size,1,1,1)
+
+    def init_bhwd(self, batch_size, height, width, device, amp):
+        self.pos_s16 = self.init_pos(batch_size, height, width, device, amp)
+
+    def forward(self, img):
+
+        img = F.avg_pool2d(img, kernel_size=2, stride=2)
+        x_8 = self.block_8_1(img)
+
+        img = F.avg_pool2d(img, kernel_size=2, stride=2)
+        x_8_2 = self.block_8_2(img)
+
+        x_8 = self.block_cat_8(torch.cat([x_8, x_8_2], dim=1))
+
+        img = F.avg_pool2d(img, kernel_size=2, stride=2)
+        x_16 = self.block_16_1(img)
+
+        x_16_2 = self.block_8_16(x_8)
+
+        x_16 = self.block_cat_16(torch.cat([x_16, x_16_2], dim=1))
+
+        x_16 = torch.cat([x_16, self.pos_s16], dim=1)
+
+        return x_16, x_8

NeuFlow/backbone_v8.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn.functional as F
+
+
+class ConvBlock(torch.nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, stride, padding):
+        super(ConvBlock, self).__init__()
+
+        self.conv1 = torch.nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode='zeros', bias=False)
+
+        self.conv2 = torch.nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False)
+
+        self.relu = torch.nn.LeakyReLU(negative_slope=0.1, inplace=False)
+
+        self.norm1 = torch.nn.BatchNorm2d(out_planes)
+
+        self.norm2 = torch.nn.BatchNorm2d(out_planes)
+
+        # self.dropout = torch.nn.Dropout(p=0.1)
+
+    def forward(self, x):
+
+        # x = self.dropout(x)
+
+        x = self.relu(self.norm1(self.conv1(x)))
+        x = self.relu(self.norm2(self.conv2(x)))
+        # x = self.relu(self.conv1(x))
+        # x = self.relu(self.conv2(x))
+
+        return x
+
+class CNNEncoder(torch.nn.Module):
+    def __init__(self, feature_dim_s16, context_dim_s16, feature_dim_s8, context_dim_s8):
+        super(CNNEncoder, self).__init__()
+
+        self.block_8_1 = ConvBlock(3, feature_dim_s8 * 2, kernel_size=8, stride=4, padding=2)
+
+        self.block_8_2 = ConvBlock(3, feature_dim_s8, kernel_size=6, stride=2, padding=2)
+
+        self.block_cat_8 = ConvBlock(feature_dim_s8 * 3, feature_dim_s8 + context_dim_s8, kernel_size=3, stride=1, padding=1)
+
+        self.block_8_16 = ConvBlock(feature_dim_s8 + context_dim_s8, feature_dim_s16 + context_dim_s16 - 2, kernel_size=6, stride=2, padding=2)
+
+    def init_pos(self, batch_size, height, width, device, amp):
+        ys, xs = torch.meshgrid(torch.arange(height, dtype=torch.half if amp else torch.float, device=device), torch.arange(width, dtype=torch.half if amp else torch.float, device=device), indexing='ij')
+        ys = (ys-height/2)
+        xs = (xs-width/2)
+        pos = torch.stack([ys, xs])
+        return pos[None].repeat(batch_size,1,1,1)
+
+    def init_bhwd(self, batch_size, height, width, device, amp):
+        self.pos_s16 = self.init_pos(batch_size, height, width, device, amp)
+
+    def forward(self, img):
+
+        img = F.avg_pool2d(img, kernel_size=2, stride=2)
+        x_8 = self.block_8_1(img)
+
+        img = F.avg_pool2d(img, kernel_size=2, stride=2)
+        x_8_2 = self.block_8_2(img)
+
+        x_8 = self.block_cat_8(torch.cat([x_8, x_8_2], dim=1))
+
+        x_16 = self.block_8_16(x_8)
+
+        x_16 = torch.cat([x_16, self.pos_s16], dim=1)
+
+        return x_16, x_8

NeuFlow/config.py

@@ -1,5 +1,7 @@
-feature_dim_s16 = 256
-hidden_dim_s16 = 128
+feature_dim_s16 = 128
+context_dim_s16 = 64
+iter_context_dim_s16 = 64
 feature_dim_s8 = 128
-hidden_dim_s8 = 96
+context_dim_s8 = 64
+iter_context_dim_s8 = 64
 feature_dim_s1 = 128

NeuFlow/neuflow.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn.functional as F
 
-from NeuFlow import backbone_v6
+from NeuFlow import backbone_v8
 from NeuFlow import transformer
 from NeuFlow import matching
 from NeuFlow import corr
@@ -17,38 +17,29 @@ class NeuFlow(torch.nn.Module):
     def __init__(self):
         super(NeuFlow, self).__init__()
 
-        self.backbone = backbone_v6.CNNEncoder(config.feature_dim_s16, config.feature_dim_s8)
+        self.backbone = backbone_v8.CNNEncoder(config.feature_dim_s16, config.context_dim_s16, config.feature_dim_s8, config.context_dim_s8)
 
-        self.cross_attn_s16 = transformer.FeatureAttention(config.feature_dim_s16, num_layers=2, ffn=True, ffn_dim_expansion=1, post_norm=True)
+        self.cross_attn_s16 = transformer.FeatureAttention(config.feature_dim_s16+config.context_dim_s16, num_layers=2, ffn=True, ffn_dim_expansion=1, post_norm=True)
 
         self.matching_s16 = matching.Matching()
 
         # self.flow_attn_s16 = transformer.FlowAttention(config.feature_dim_s16)
-        
-        self.merge_s8 = torch.nn.Sequential(torch.nn.Conv2d(config.feature_dim_s16 + config.feature_dim_s8 + 2, config.feature_dim_s8, kernel_size=3, stride=1, padding=1, bias=False),
-                                              torch.nn.GELU(),
-                                              torch.nn.Conv2d(config.feature_dim_s8, config.feature_dim_s8, kernel_size=3, stride=1, padding=1, bias=False),
-                                              torch.nn.Tanh())
 
         self.corr_block_s16 = corr.CorrBlock(radius=4, levels=1)
         self.corr_block_s8 = corr.CorrBlock(radius=4, levels=1)
+        
+        self.merge_s8 = torch.nn.Sequential(torch.nn.Conv2d(config.feature_dim_s16 + config.feature_dim_s8, config.feature_dim_s8, kernel_size=3, stride=1, padding=1, bias=False),
+                                              torch.nn.GELU(),
+                                              torch.nn.Conv2d(config.feature_dim_s8, config.feature_dim_s8, kernel_size=3, stride=1, padding=1, bias=False))
 
-        self.context_s16 = torch.nn.Sequential(torch.nn.Conv2d(config.feature_dim_s16, config.hidden_dim_s16, kernel_size=3, stride=1, padding=1, bias=False),
-                                            torch.nn.GELU(),
-                                            torch.nn.Conv2d(config.hidden_dim_s16, config.hidden_dim_s16, kernel_size=3, stride=1, padding=1, bias=False))
-
-        self.context_merge_s8 = torch.nn.Sequential(torch.nn.Conv2d(config.hidden_dim_s16 + config.feature_dim_s8, config.hidden_dim_s8, kernel_size=3, stride=1, padding=1, bias=False),
+        self.context_merge_s8 = torch.nn.Sequential(torch.nn.Conv2d(config.context_dim_s16 + config.context_dim_s8, config.context_dim_s8, kernel_size=3, stride=1, padding=1, bias=False),
                                            torch.nn.GELU(),
-                                           torch.nn.Conv2d(config.hidden_dim_s8, config.hidden_dim_s8, kernel_size=3, stride=1, padding=1, bias=False),
-                                           torch.nn.Tanh())
-
-        self.refine_s16 = refine.Refine(config.hidden_dim_s16, num_layers=6, levels=1, radius=4)
-        self.refine_s8 = refine.Refine(config.hidden_dim_s8, num_layers=6, levels=1, radius=4)
+                                           torch.nn.Conv2d(config.context_dim_s8, config.context_dim_s8, kernel_size=3, stride=1, padding=1, bias=False))
 
-        # self.conv_s16 = backbone_v6.ConvBlock(3, config.feature_dim_s1 * 2, kernel_size=16, stride=16, padding=0)
-        # self.upsample_s16 = upsample.UpSample(config.feature_dim_s1 * 2, upsample_factor=16)
+        self.refine_s16 = refine.Refine(config.context_dim_s16, config.iter_context_dim_s16, num_layers=5, levels=1, radius=4, inter_dim=128)
+        self.refine_s8 = refine.Refine(config.context_dim_s8, config.iter_context_dim_s8, num_layers=5, levels=1, radius=4, inter_dim=96)
 
-        self.conv_s8 = backbone_v6.ConvBlock(3, config.feature_dim_s1, kernel_size=8, stride=8, padding=0)
+        self.conv_s8 = backbone_v8.ConvBlock(3, config.feature_dim_s1, kernel_size=8, stride=8, padding=0)
         self.upsample_s8 = upsample.UpSample(config.feature_dim_s1, upsample_factor=8)
 
         for p in self.parameters():
@@ -57,7 +48,7 @@ def __init__(self):
 
     def init_bhwd(self, batch_size, height, width, device, amp=True):
 
-        self.backbone.init_bhwd(batch_size*2, height, width, device, amp)
+        self.backbone.init_bhwd(batch_size*2, height//16, width//16, device, amp)
 
         self.matching_s16.init_bhwd(batch_size, height//16, width//16, device, amp)
 
@@ -67,7 +58,19 @@ def init_bhwd(self, batch_size, height, width, device, amp=True):
         self.refine_s16.init_bhwd(batch_size, height//16, width//16, device, amp)
         self.refine_s8.init_bhwd(batch_size, height//8, width//8, device, amp)
 
-    def forward(self, img0, img1, iters_s16=1, iters_s8=6):
+        self.init_iter_context_s16 = torch.zeros(batch_size, config.iter_context_dim_s16, height//16, width//16, device=device, dtype=torch.half if amp else torch.float)
+        self.init_iter_context_s8 = torch.zeros(batch_size, config.iter_context_dim_s8, height//8, width//8, device=device, dtype=torch.half if amp else torch.float)
+
+    def split_features(self, features, context_dim, feature_dim):
+
+        context, features = torch.split(features, [context_dim, feature_dim], dim=1)
+
+        context, _ = context.chunk(chunks=2, dim=0)
+        feature0, feature1 = features.chunk(chunks=2, dim=0)
+
+        return features, torch.relu(context)
+
+    def forward(self, img0, img1, iters_s16=3, iters_s8=7):
 
         flow_list = []
 
@@ -78,6 +81,9 @@ def forward(self, img0, img1, iters_s16=1, iters_s8=6):
 
         features_s16 = self.cross_attn_s16(features_s16)
 
+        features_s16, context_s16 = self.split_features(features_s16, config.context_dim_s16, config.feature_dim_s16)
+        features_s8, context_s8 = self.split_features(features_s8, config.context_dim_s8, config.feature_dim_s8)
+
         feature0_s16, feature1_s16 = features_s16.chunk(chunks=2, dim=0)
 
         flow0 = self.matching_s16.global_correlation_softmax(feature0_s16, feature1_s16)
@@ -86,18 +92,17 @@ def forward(self, img0, img1, iters_s16=1, iters_s8=6):
 
         corr_pyr_s16 = self.corr_block_s16.init_corr_pyr(feature0_s16, feature1_s16)
 
-        context_s16 = self.context_s16(feature0_s16)
-        iter_context_s16 = context_s16.clone()
+        iter_context_s16 = self.init_iter_context_s16
 
         for i in range(iters_s16):
 
             if self.training and i > 0:
                 flow0 = flow0.detach()
-                # iter_feature0_s16 = iter_feature0_s16.detach()
+                # iter_context_s16 = iter_context_s16.detach()
 
             corrs = self.corr_block_s16(corr_pyr_s16, flow0)
 
-            iter_context_s16, delta_flow = self.refine_s16(corrs, iter_context_s16, flow0)
+            iter_context_s16, delta_flow = self.refine_s16(corrs, context_s16, iter_context_s16, flow0)
 
             flow0 = flow0 + delta_flow
 
@@ -117,18 +122,20 @@ def forward(self, img0, img1, iters_s16=1, iters_s8=6):
 
         context_s16 = F.interpolate(context_s16, scale_factor=2, mode='nearest')
 
-        context_s8 = self.context_merge_s8(torch.cat([feature0_s8, context_s16], dim=1))
-        iter_context_s8 = context_s8.clone()
+        context_s8 = torch.zeros_like(context_s8)
+        context_s8 = self.context_merge_s8(torch.cat([context_s8, context_s16], dim=1))
+
+        iter_context_s8 = self.init_iter_context_s8
 
         for i in range(iters_s8):
 
             if self.training and i > 0:
                 flow0 = flow0.detach()
-                # iter_feature0_s8 = iter_feature0_s8.detach()
+                # iter_context_s8 = iter_context_s8.detach()
 
             corrs = self.corr_block_s8(corr_pyr_s8, flow0)
 
-            iter_context_s8, delta_flow = self.refine_s8(corrs, iter_context_s8, flow0)
+            iter_context_s8, delta_flow = self.refine_s8(corrs, context_s8, iter_context_s8, flow0)
 
             flow0 = flow0 + delta_flow
 

NeuFlow/refine.py

@@ -1,41 +1,49 @@
 import torch
+from NeuFlow import utils
 
 
 class ConvBlock(torch.nn.Module):
     def __init__(self, in_planes, out_planes, kernel_size, stride, padding):
         super(ConvBlock, self).__init__()
 
-        self.conv = torch.nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode='zeros', bias=True)
+        self.conv = torch.nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode='zeros', bias=False)
         self.relu = torch.nn.LeakyReLU(negative_slope=0.1, inplace=False)
 
     def forward(self, x):
         return self.relu(self.conv(x))
 
 class Refine(torch.nn.Module):
-    def __init__(self, feature_dim, num_layers, levels, radius):
+    def __init__(self, context_dim, iter_context_dim, num_layers, levels, radius, inter_dim):
         super(Refine, self).__init__()
 
         self.radius = radius
 
-        self.conv1 = ConvBlock((radius*2+1)**2*levels+feature_dim+2+1, feature_dim, kernel_size=3, stride=1, padding=1)
+        self.conv1 = ConvBlock((radius*2+1)**2*levels+context_dim+iter_context_dim+2+1, context_dim+iter_context_dim, kernel_size=3, stride=1, padding=1)
 
-        self.conv_layers = torch.nn.ModuleList([ConvBlock(feature_dim, feature_dim, kernel_size=3, stride=1, padding=1)
+        self.conv2 = ConvBlock(context_dim+iter_context_dim, inter_dim, kernel_size=3, stride=1, padding=1)
+
+        self.conv_layers = torch.nn.ModuleList([ConvBlock(inter_dim, inter_dim, kernel_size=3, stride=1, padding=1)
                                                 for i in range(num_layers)])
 
-        self.conv2 = torch.nn.Conv2d(feature_dim, feature_dim+2, kernel_size=3, stride=1, padding=1, padding_mode='zeros', bias=True)
+        self.conv3 = torch.nn.Conv2d(inter_dim, iter_context_dim+2, kernel_size=3, stride=1, padding=1, padding_mode='zeros', bias=True)
+
+        self.hidden_act = torch.nn.Tanh()
+        # self.hidden_norm = torch.nn.BatchNorm2d(feature_dim)
 
     def init_bhwd(self, batch_size, height, width, device, amp):
         self.radius_emb = torch.tensor(self.radius, dtype=torch.half if amp else torch.float, device=device).view(1,-1,1,1).expand([batch_size,1,height,width])
 
-    def forward(self, corrs, feature0, flow0):
+    def forward(self, corrs, context, iter_context, flow0):
 
-        x = torch.cat([corrs, feature0, flow0, self.radius_emb], dim=1)
+        x = torch.cat([corrs, context, iter_context, flow0, self.radius_emb], dim=1)
 
         x = self.conv1(x)
 
+        x = self.conv2(x)
+
         for layer in self.conv_layers:
             x = layer(x)
 
-        x = self.conv2(x)
+        x = self.conv3(x)
 
-        return torch.tanh(x[:,2:]), x[:,:2]
+        return self.hidden_act(x[:,2:]), x[:,:2]

NeuFlow/utils.py

@@ -1,8 +1,9 @@
 import torch
 import torch.nn.functional as F
 
-def normalize_img(img, mean, std):
-    return (img / 255. - mean) / std
+# def normalize(x):
+#     x_min = x.min()
+#     return (x - x_min) / (x.max() - x_min)
 
 def coords_grid(b, h, w, device, amp):
     ys, xs = torch.meshgrid(torch.arange(h, dtype=torch.half if amp else torch.float, device=device), torch.arange(w, dtype=torch.half if amp else torch.float, device=device), indexing='ij')  # [H, W]