compatibility changes

IGEV-Stereo/core/geometry.py

@@ -60,7 +60,6 @@ class Combined_Geo_Encoding_Volume:
     
     @staticmethod
     def corr(fmap1, fmap2):
-        # batch, dim, ht, wd
         B, D, H, W1 = fmap1.shape
         _, _, _, W2 = fmap2.shape
         fmap1 = fmap1.view(B, D, H, W1)

IGEV-Stereo/core/igev_stereo.py

@@ -167,7 +167,6 @@ class IGEVStereo(nn.Module):
             match_right = self.desc(self.conv(features_right[0]))
             gwc_volume = build_gwc_volume(match_left, match_right, self.args.max_disp//4, 8)
             gwc_volume = self.corr_stem(gwc_volume)
-            # 3d unet
             gwc_volume = self.corr_feature_att(gwc_volume, features_left[0])
             geo_encoding_volume = self.cost_agg(gwc_volume, features_left)
 

IGEV-Stereo/core/stereo_datasets.py

@@ -73,6 +73,7 @@ class StereoDataset(data.Dataset):
         img2 = np.array(img2).astype(np.uint8)
 
         disp = np.array(disp).astype(np.float32)
+        assert not (True in np.isnan(disp))
 
         flow = np.stack([disp, np.zeros_like(disp)], axis=-1)
 
@@ -293,7 +294,7 @@ class CREStereo(StereoDataset):
 
 
 class Middlebury(StereoDataset):
-    def __init__(self, aug_params=None, root='/data/Middlebury', split='H'):
+    def __init__(self, aug_params=None, root='/data/Middlebury', split='F'):
         super(Middlebury, self).__init__(aug_params, sparse=True, reader=frame_utils.readDispMiddlebury)
         assert os.path.exists(root)
         assert split in "FHQ"
@@ -352,7 +353,7 @@ def fetch_dataloader(args):
         train_dataset = new_dataset if train_dataset is None else train_dataset + new_dataset
 
     train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size, 
-        pin_memory=True, shuffle=True, num_workers=12, drop_last=True)
+        pin_memory=True, shuffle=True, num_workers=int(os.environ.get('SLURM_CPUS_PER_TASK', 6))-2, drop_last=True)
 
     logging.info('Training with %d image pairs' % len(train_dataset))
     return train_loader

IGEV-Stereo/demo_imgs.py

@@ -15,6 +15,7 @@ from PIL import Image
 from matplotlib import pyplot as plt
 import os
 import cv2
+from torch.profiler import profile, record_function, ProfilerActivity
 
 def load_image(imfile):
     img = np.array(Image.open(imfile)).astype(np.uint8)
@@ -44,8 +45,14 @@ def demo(args):
 
             padder = InputPadder(image1.shape, divis_by=32)
             image1, image2 = padder.pad(image1, image2)
-
-            disp = model(image1, image2, iters=args.valid_iters, test_mode=True)
+            with profile(
+                    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+                    with_stack=True,
+            ) as prof:
+                disp = model(image1, image2, iters=args.valid_iters, test_mode=True)
+            print(prof.key_averages(group_by_stack_n=5).table(sort_by="self_cuda_time_total", row_limit=10))
+            prof.export_chrome_trace("./trace.json")
+            # disp = model(image1, image2, iters=args.valid_iters, test_mode=True)
             disp = disp.cpu().numpy()
             disp = padder.unpad(disp)
             file_stem = imfile1.split('/')[-2]
@@ -81,6 +88,7 @@ if __name__ == '__main__':
     parser.add_argument('--slow_fast_gru', action='store_true', help="iterate the low-res GRUs more frequently")
     parser.add_argument('--n_gru_layers', type=int, default=3, help="number of hidden GRU levels")
     parser.add_argument('--max_disp', type=int, default=192, help="max disp of geometry encoding volume")
+    parser.add_argument('--freeze_backbone_params', action="store_true", help="freeze backbone parameters")
 
     args = parser.parse_args()
 

IGEV-Stereo/evaluate_stereo.py

@@ -169,7 +169,7 @@ def validate_sceneflow(model, iters=32, mixed_prec=False, max_disp=192):
 
 
 @torch.no_grad()
-def validate_middlebury(model, iters=32, split='H', mixed_prec=False, max_disp=192):
+def validate_middlebury(model, iters=32, split='F', mixed_prec=False, max_disp=192):
     """ Peform validation using the Middlebury-V3 dataset """
     model.eval()
     aug_params = {}

IGEV-Stereo/train_stereo.py

@@ -22,6 +22,7 @@ from evaluate_stereo import *
 import core.stereo_datasets as datasets
 import torch.nn.functional as F
 
+torch.backends.cudnn.benchmark = True
 
 ckpt_path = './checkpoints/igev_stereo'
 log_path = './checkpoints/igev_stereo'