remove cuda correlation

visinf · Mar 14, 2022 · dacd07b · dacd07b
1 parent 4d7f6aa
commit dacd07b
Show file tree

Hide file tree

Showing 10 changed files with 50 additions and 46 deletions.
diff --git a/datasets/kitti_combined.py b/datasets/kitti_combined.py
@@ -72,7 +72,7 @@ def __init__(self,
         if images_root_2015 is not None:
 
             if not os.path.isdir(images_root_2015):
-                raise ValueError("Image directory '%s' not found!")
+                raise ValueError("Image directory not found! {}".format(images_root_2015))
 
             all_img1_2015_filenames = sorted(glob(os.path.join(images_root_2015, "*_10.png")))
             all_img2_2015_filenames = sorted(glob(os.path.join(images_root_2015, "*_11.png")))
@@ -86,7 +86,7 @@ def __init__(self,
         if images_root_2012 is not None:
 
             if not os.path.isdir(images_root_2012):
-                raise ValueError("Image directory '%s' not found!")
+                raise ValueError("Image directory not found! {}".format(images_root_2012))
 
             all_img1_2012_filenames = sorted(glob(os.path.join(images_root_2012, "*_10.png")))
             all_img2_2012_filenames = sorted(glob(os.path.join(images_root_2012, "*_11.png")))
@@ -205,9 +205,9 @@ def __init__(self,
         if images_root_2015 is not None and flow_root_2015 is not None:
 
             if not os.path.isdir(images_root_2015):
-                raise ValueError("Image directory '%s' not found!")
+                raise ValueError("Image directory not found!  {}".format(images_root_2015))
             if not os.path.isdir(flow_root_2015):
-                raise ValueError("Flow directory '%s' not found!")
+                raise ValueError("Flow directory not found!  {}".format(flow_root_2015))
 
             all_img1_2015_filenames = sorted(glob(os.path.join(images_root_2015, "*_10.png")))
             all_img2_2015_filenames = sorted(glob(os.path.join(images_root_2015, "*_11.png")))            

diff --git a/models/IRR_PWC.py b/models/IRR_PWC.py
@@ -6,7 +6,6 @@
 from .pwc_modules import conv, upsample2d_as, rescale_flow, initialize_msra, compute_cost_volume
 from .pwc_modules import WarpingLayer, FeatureExtractor, ContextNetwork, FlowEstimatorDense, OccContextNetwork, OccEstimatorDense
 from .irr_modules import OccUpsampleNetwork, RefineFlow, RefineOcc
-# from .correlation_package.correlation import Correlation
 import copy
 
 
@@ -88,8 +87,6 @@ def forward(self, input_dict):
                     x1_warp = self.warping_layer(x1, flow_b, height_im, width_im, self._div_flow)
 
                 # correlation
-                # out_corr_f = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x1, x2_warp)
-                # out_corr_b = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x2, x1_warp)
                 out_corr_f = compute_cost_volume(x1, x2_warp, self.corr_params)
                 out_corr_b = compute_cost_volume(x2, x1_warp, self.corr_params)
 

diff --git a/models/pwcnet.py b/models/pwcnet.py
@@ -3,9 +3,8 @@
 import torch
 import torch.nn as nn
 
-from .pwc_modules import upsample2d_as, initialize_msra
+from .pwc_modules import upsample2d_as, initialize_msra, compute_cost_volume
 from .pwc_modules import WarpingLayer, FeatureExtractor, ContextNetwork, FlowEstimatorDense
-from .correlation_package.correlation import Correlation
 
 class PWCNet(nn.Module):
     def __init__(self, args, div_flow=0.05):
@@ -36,7 +35,9 @@ def __init__(self, args, div_flow=0.05):
             self.flow_estimators.append(layer)
 
         self.context_networks = ContextNetwork(self.dim_corr + 32 + 2 + 448 + 2)
-
+
+        self.corr_params = {"pad_size": self.search_range, "kernel_size": 1, "max_disp": self.search_range, "stride1": 1, "stride2": 1, "corr_multiply": 1}
+
         initialize_msra(self.modules())
 
     def forward(self, input_dict):
@@ -69,7 +70,7 @@ def forward(self, input_dict):
                 x2_warp = self.warping_layer(x2, flow, height_im, width_im, self._div_flow)
 
             # correlation
-            out_corr = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x1, x2_warp)
+            out_corr = compute_cost_volume(x1, x2_warp, self.corr_params)
             out_corr_relu = self.leakyRELU(out_corr)
 
             # flow estimator

diff --git a/models/pwcnet_bi.py b/models/pwcnet_bi.py
@@ -3,9 +3,8 @@
 import torch
 import torch.nn as nn
 
-from .pwc_modules import upsample2d_as, initialize_msra
+from .pwc_modules import upsample2d_as, initialize_msra, compute_cost_volume
 from .pwc_modules import WarpingLayer, FeatureExtractor, ContextNetwork, FlowEstimatorDense
-from .correlation_package.correlation import Correlation
 
 class PWCNet(nn.Module):
     def __init__(self, args, div_flow=0.05):
@@ -36,7 +35,8 @@ def __init__(self, args, div_flow=0.05):
             self.flow_estimators.append(layer)
 
         self.context_networks = ContextNetwork(self.dim_corr + 32 + 2 + 448 + 2)
-
+        self.corr_params = {"pad_size": self.search_range, "kernel_size": 1, "max_disp": self.search_range, "stride1": 1, "stride2": 1, "corr_multiply": 1}
+
         initialize_msra(self.modules())
 
     def forward(self, input_dict):
@@ -72,10 +72,10 @@ def forward(self, input_dict):
                 x2_warp = self.warping_layer(x2, flow_f, height_im, width_im, self._div_flow)
                 x1_warp = self.warping_layer(x1, flow_b, height_im, width_im, self._div_flow)
 
-            # correlation            
-            out_corr_f = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x1, x2_warp)
-            out_corr_b = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x2, x1_warp)
-
+            # correlation
+            out_corr_f = compute_cost_volume(x1, x2_warp, self.corr_params)
+            out_corr_b = compute_cost_volume(x2, x1_warp, self.corr_params)
+                
             out_corr_relu_f = self.leakyRELU(out_corr_f)
             out_corr_relu_b = self.leakyRELU(out_corr_b)
 

diff --git a/models/pwcnet_irr.py b/models/pwcnet_irr.py
@@ -3,9 +3,8 @@
 import torch
 import torch.nn as nn
 
-from .pwc_modules import conv, rescale_flow, upsample2d_as, initialize_msra
+from .pwc_modules import conv, rescale_flow, upsample2d_as, initialize_msra, compute_cost_volume
 from .pwc_modules import WarpingLayer, FeatureExtractor, ContextNetwork, FlowEstimatorDense
-from .correlation_package.correlation import Correlation
 
 class PWCNet(nn.Module):
     def __init__(self, args, div_flow=0.05):
@@ -33,7 +32,9 @@ def __init__(self, args, div_flow=0.05):
                                        conv(96, 32, kernel_size=1, stride=1, dilation=1),
                                        conv(64, 32, kernel_size=1, stride=1, dilation=1),
                                        conv(32, 32, kernel_size=1, stride=1, dilation=1)])
-
+
+        self.corr_params = {"pad_size": self.search_range, "kernel_size": 1, "max_disp": self.search_range, "stride1": 1, "stride2": 1, "corr_multiply": 1}
+
         initialize_msra(self.modules())
 
     def forward(self, input_dict):
@@ -66,7 +67,7 @@ def forward(self, input_dict):
                 x2_warp = self.warping_layer(x2, flow, height_im, width_im, self._div_flow)
 
             # correlation
-            out_corr = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x1, x2_warp)
+            out_corr = compute_cost_volume(x1, x2_warp, self.corr_params)
             out_corr_relu = self.leakyRELU(out_corr)
 
             # concat and estimate flow

diff --git a/models/pwcnet_irr_bi.py b/models/pwcnet_irr_bi.py
@@ -3,9 +3,8 @@
 import torch
 import torch.nn as nn
 
-from .pwc_modules import conv, rescale_flow, upsample2d_as, initialize_msra
+from .pwc_modules import conv, rescale_flow, upsample2d_as, initialize_msra, compute_cost_volume
 from .pwc_modules import WarpingLayer, FeatureExtractor, ContextNetwork, FlowEstimatorDense
-from .correlation_package.correlation import Correlation
 
 class PWCNet(nn.Module):
     def __init__(self, args, div_flow=0.05):
@@ -33,7 +32,9 @@ def __init__(self, args, div_flow=0.05):
                                        conv(96, 32, kernel_size=1, stride=1, dilation=1),
                                        conv(64, 32, kernel_size=1, stride=1, dilation=1),
                                        conv(32, 32, kernel_size=1, stride=1, dilation=1)])
-
+
+        self.corr_params = {"pad_size": self.search_range, "kernel_size": 1, "max_disp": self.search_range, "stride1": 1, "stride2": 1, "corr_multiply": 1}
+
         initialize_msra(self.modules())
 
     def forward(self, input_dict):
@@ -70,8 +71,8 @@ def forward(self, input_dict):
                 x1_warp = self.warping_layer(x1, flow_b, height_im, width_im, self._div_flow)
 
             # correlation
-            out_corr_f = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x1, x2_warp)
-            out_corr_b = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x2, x1_warp)
+            out_corr_f = compute_cost_volume(x1, x2_warp, self.corr_params)
+            out_corr_b = compute_cost_volume(x2, x1_warp, self.corr_params)
             out_corr_relu_f = self.leakyRELU(out_corr_f)
             out_corr_relu_b = self.leakyRELU(out_corr_b)
 

diff --git a/models/pwcnet_irr_occ.py b/models/pwcnet_irr_occ.py
@@ -3,9 +3,8 @@
 import torch
 import torch.nn as nn
 
-from .pwc_modules import conv, rescale_flow, upsample2d_as, initialize_msra
+from .pwc_modules import conv, rescale_flow, upsample2d_as, initialize_msra, compute_cost_volume
 from .pwc_modules import WarpingLayer, FeatureExtractor, FlowEstimatorDense, ContextNetwork, OccEstimatorDense, OccContextNetwork
-from .correlation_package.correlation import Correlation
 
 class PWCNet(nn.Module):
     def __init__(self, args, div_flow=0.05):
@@ -36,7 +35,9 @@ def __init__(self, args, div_flow=0.05):
                                        conv(96, 32, kernel_size=1, stride=1, dilation=1),
                                        conv(64, 32, kernel_size=1, stride=1, dilation=1),
                                        conv(32, 32, kernel_size=1, stride=1, dilation=1)])
-
+
+        self.corr_params = {"pad_size": self.search_range, "kernel_size": 1, "max_disp": self.search_range, "stride1": 1, "stride2": 1, "corr_multiply": 1}
+
         initialize_msra(self.modules())
 
     def forward(self, input_dict):
@@ -72,7 +73,7 @@ def forward(self, input_dict):
                 x2_warp = self.warping_layer(x2, flow, height_im, width_im, self._div_flow)
 
             # correlation
-            out_corr = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x1, x2_warp)
+            out_corr = compute_cost_volume(x1, x2_warp, self.corr_params)
             out_corr_relu = self.leakyRELU(out_corr)
 
             # concat and estimate flow

diff --git a/models/pwcnet_irr_occ_bi.py b/models/pwcnet_irr_occ_bi.py
@@ -3,9 +3,8 @@
 import torch
 import torch.nn as nn
 
-from .pwc_modules import conv, rescale_flow, upsample2d_as, initialize_msra
+from .pwc_modules import conv, rescale_flow, upsample2d_as, initialize_msra, compute_cost_volume
 from .pwc_modules import WarpingLayer, FeatureExtractor, FlowEstimatorDense, ContextNetwork, OccEstimatorDense, OccContextNetwork
-from .correlation_package.correlation import Correlation
 
 class PWCNet(nn.Module):
     def __init__(self, args, div_flow=0.05):
@@ -36,7 +35,9 @@ def __init__(self, args, div_flow=0.05):
                                        conv(96, 32, kernel_size=1, stride=1, dilation=1),
                                        conv(64, 32, kernel_size=1, stride=1, dilation=1),
                                        conv(32, 32, kernel_size=1, stride=1, dilation=1)])
-
+
+        self.corr_params = {"pad_size": self.search_range, "kernel_size": 1, "max_disp": self.search_range, "stride1": 1, "stride2": 1, "corr_multiply": 1}
+
         initialize_msra(self.modules())
 
     def forward(self, input_dict):
@@ -78,8 +79,8 @@ def forward(self, input_dict):
                 x1_warp = self.warping_layer(x1, flow_b, height_im, width_im, self._div_flow)
 
             # correlation
-            out_corr_f = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x1, x2_warp)
-            out_corr_b = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x2, x1_warp)
+            out_corr_f = compute_cost_volume(x1, x2_warp, self.corr_params)
+            out_corr_b = compute_cost_volume(x2, x1_warp, self.corr_params)
             out_corr_relu_f = self.leakyRELU(out_corr_f)
             out_corr_relu_b = self.leakyRELU(out_corr_b)
 

diff --git a/models/pwcnet_occ.py b/models/pwcnet_occ.py
@@ -3,9 +3,8 @@
 import torch
 import torch.nn as nn
 
-from .pwc_modules import upsample2d_as, initialize_msra
+from .pwc_modules import upsample2d_as, initialize_msra, compute_cost_volume
 from .pwc_modules import WarpingLayer, FeatureExtractor, FlowEstimatorDense, ContextNetwork, OccEstimatorDense, OccContextNetwork
-from .correlation_package.correlation import Correlation
 
 class PWCNet(nn.Module):
     def __init__(self, args, div_flow=0.05):
@@ -42,7 +41,9 @@ def __init__(self, args, div_flow=0.05):
 
         self.context_networks = ContextNetwork(self.dim_corr + 32 + 2 + 448 + 2)
         self.context_networks_occ = OccContextNetwork(self.dim_corr + 32 + 1 + 448 + 1)
-
+
+        self.corr_params = {"pad_size": self.search_range, "kernel_size": 1, "max_disp": self.search_range, "stride1": 1, "stride2": 1, "corr_multiply": 1}
+
         initialize_msra(self.modules())
 
     def forward(self, input_dict):
@@ -77,8 +78,8 @@ def forward(self, input_dict):
                 occ = upsample2d_as(occ, x1, mode="bilinear")
                 x2_warp = self.warping_layer(x2, flow, height_im, width_im, self._div_flow)
 
-            # correlation
-            out_corr = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x1, x2_warp)            
+            # correlation  
+            out_corr = compute_cost_volume(x1, x2_warp, self.corr_params)
             out_corr_relu = self.leakyRELU(out_corr)
 
             # flow estimator

diff --git a/models/pwcnet_occ_bi.py b/models/pwcnet_occ_bi.py
@@ -3,9 +3,8 @@
 import torch
 import torch.nn as nn
 
-from .pwc_modules import upsample2d_as, initialize_msra
+from .pwc_modules import upsample2d_as, initialize_msra, compute_cost_volume
 from .pwc_modules import WarpingLayer, FeatureExtractor, FlowEstimatorDense, ContextNetwork, OccEstimatorDense, OccContextNetwork
-from .correlation_package.correlation import Correlation
 
 class PWCNet(nn.Module):
     def __init__(self, args, div_flow=0.05):
@@ -42,7 +41,9 @@ def __init__(self, args, div_flow=0.05):
 
         self.context_networks = ContextNetwork(self.dim_corr + 32 + 2 + 448 + 2)
         self.context_networks_occ = OccContextNetwork(self.dim_corr + 32 + 1 + 448 + 1)
-
+
+        self.corr_params = {"pad_size": self.search_range, "kernel_size": 1, "max_disp": self.search_range, "stride1": 1, "stride2": 1, "corr_multiply": 1}
+
         initialize_msra(self.modules())
 
     def forward(self, input_dict):
@@ -84,8 +85,8 @@ def forward(self, input_dict):
                 x1_warp = self.warping_layer(x1, flow_b, height_im, width_im, self._div_flow)
 
             # correlation
-            out_corr_f = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x1, x2_warp)
-            out_corr_b = Correlation(pad_size=self.search_range, kernel_size=1, max_displacement=self.search_range, stride1=1, stride2=1, corr_multiply=1)(x2, x1_warp)            
+            out_corr_f = compute_cost_volume(x1, x2_warp, self.corr_params)
+            out_corr_b = compute_cost_volume(x2, x1_warp, self.corr_params)
             out_corr_relu_f = self.leakyRELU(out_corr_f)
             out_corr_relu_b = self.leakyRELU(out_corr_b)