IGEV-Stereo/core/stereo_datasets.py

@@ -353,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/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'