Speed up dataloader and set batch size 64

params_storage/params_fig10b_single_gpu.py

@@ -32,7 +32,7 @@ opt_exp.gpu_ids = ['0']           # Single GPU
 opt_exp.data = "rw_to_rw"        # Jacobs Hall complex env
 opt_exp.n_decoders = 2            # Encoder + location decoder + consistency decoder
 
-opt_exp.batch_size = 8            # Reduced from 32 for 6 GB VRAM
+opt_exp.batch_size = 64           # Larger batch for H100-class GPUs
 opt_exp.ds_step_trn = 1
 opt_exp.ds_step_tst = 1
 opt_exp.weight_decay = 1e-5

params_storage/params_mamba_single_gpu.py

@@ -30,7 +30,7 @@ opt_exp.data = "rw_to_rw"           # Same Jacobs Hall dataset as Fig 10b
 opt_exp.n_decoders = 1               # Single decoder (location only)
 opt_exp.model_type = "mamba"         # Selects Mamba pipeline in train_and_test.py
 
-opt_exp.batch_size = 8
+opt_exp.batch_size = 64
 opt_exp.ds_step_trn = 1
 opt_exp.ds_step_tst = 1
 opt_exp.weight_decay = 1e-5

train_and_test.py

@@ -121,14 +121,26 @@ Samples are read on demand from HDF5 files - uses ~0 RAM.
 print('Indexing training files (lazy, no data loaded into RAM)...')
 train_data = LazyMultiHDF5Dataset(trainpath)
 train_loader = torch.utils.data.DataLoader(
-    train_data, batch_size=opt_exp.batch_size, shuffle=True, num_workers=0)
+    train_data,
+    batch_size=opt_exp.batch_size,
+    shuffle=True,
+    num_workers=4,
+    pin_memory=True,
+    persistent_workers=True,
+)
 print(f'# training samples = {len(train_data)}')
 print(f'# training mini batch = {len(train_loader)}')
 
 print('Indexing testing files...')
 test_data = LazyMultiHDF5Dataset(testpath)
 test_loader = torch.utils.data.DataLoader(
-    test_data, batch_size=opt_exp.batch_size, shuffle=False, num_workers=0)
+    test_data,
+    batch_size=opt_exp.batch_size,
+    shuffle=False,
+    num_workers=4,
+    pin_memory=True,
+    persistent_workers=True,
+)
 print(f'# testing samples = {len(test_data)}')
 print(f'# testing mini batch = {len(test_loader)}')
 print('Data Loaded (lazy)')