小bug排除，大bug还得麻烦师兄🌹🌹🌹

configs/baseline.yaml

@@ -1,11 +1,10 @@
 timesteps: [0,201]
 exp_setting: 'intra-subject'
-eeg_backbone: 'EEGProjectLayerv2'
+eeg_backbone: 'EEGProject'
 vision_backbone: 'ViT-bigG-14'
 info: ''
 name: "${exp_setting}_${eeg_backbone}_${vision_backbone}"
-seed: 0
-
+seed: 97
 
 models:
   brain:
@@ -16,13 +15,14 @@ models:
         timesteps: ${timesteps}
 
 data:
-  subjects: ['sub-10']
+  subjects: ['sub-01']
   selected_ch: ['P7', 'P5', 'P3', 'P1','Pz', 'P2', 'P4', 'P6', 'P8', 'PO7', 'PO3', 'POz', 'PO4', 'PO8','O1', 'Oz', 'O2']
   model_type: ${vision_backbone}
   train_batch_size: 1024
   val_batch_size: 200
   test_batch_size: 200
   uncertainty_aware: False
+  single_uncertainty_aware: False
   mixco: True
   blur_type:
     target: inpating_data.DirectT

configs/baseline_ubp.yaml

@@ -1,23 +1,22 @@
-timesteps: [25,250]
+timesteps: [250,600]
 blur_kernel_size: 51
 system_g: 3
 c: 10
 
 exp_setting: 'intra-subject'
 eeg_backbone: 'Ours'
-vision_backbone: 'ViT-H-14'
-info: '-ubp0.05-[25,250]'
+vision_backbone: 'ViT-bigG-14'
+info: ''
 name: "${exp_setting}_${eeg_backbone}_${vision_backbone}"
-seed: 0
-alpha: 0.05
+seed: 97
 
 
 models:
   brain:
     target: "eeg_backbone.${eeg_backbone}"
     params:
         c_num: 17
-        z_dim: 1024
+        z_dim: 1280
         timesteps: ${timesteps}
 
 data:
@@ -48,7 +47,5 @@ train:
   epoch: 50
   optimizer: AdamW
   lr: 1e-4
-  scheduler: none
-  min_lr: 1e-6
 
-save_dir: 'exp'
+save_dir: 'exp'

configs/baseline_ubp_GX.yaml

@@ -8,7 +8,7 @@ eeg_backbone: 'Ours_GX'
 vision_backbone: 'ViT-H-14'
 info: '-ubp-nhead[1024, 1, 0, 0]'
 name: "${exp_setting}_${eeg_backbone}_${vision_backbone}"
-seed: 0
+seed: 97
 
 
 models:

configs/baseline_ubp_sdxl.yaml

@@ -8,7 +8,7 @@ eeg_backbone: 'Ours_bn_remake'
 vision_backbone: 'ViT-bigG-14'
 info: ''
 name: "${exp_setting}_${eeg_backbone}_${vision_backbone}"
-seed: 0
+seed: 97
 
 
 models:

configs/baseline_ubp_sota.yaml

@@ -8,7 +8,7 @@ eeg_backbone: 'Ours'
 vision_backbone: 'ViT-H-14'
 info: '-ubp0.05-[25,250]-lr3e-4'
 name: "${exp_setting}_${eeg_backbone}_${vision_backbone}"
-seed: 0
+seed: 97
 alpha: 0.05
 
 
@@ -52,3 +52,4 @@ train:
   min_lr: 1e-6
 
 save_dir: 'exp'
+

configs/baseline_ubp_vae.yaml

@@ -8,7 +8,7 @@ eeg_backbone: 'EEGProjectLayerv2'
 vision_backbone: 'ViT-bigG-14'
 info: ''
 name: "${exp_setting}_${eeg_backbone}_${vision_backbone}"
-seed: 0
+seed: 97
 
 
 models:

src/retrieval/data_eeg.py

@@ -111,7 +111,7 @@ def load_eeg_data(config, exp_setting='intra-subject'):
 class EEGDataset(Dataset):
     def __init__(self, config, mode):
         self.config = config
-        self.data_dir = './data/EEG_ViRe_1202_cov_norm'
+        self.data_dir = './data/EEG_ViRe_1120'
         self.subjects = config['data']['subjects']
         print(f'subjects:{self.subjects}')
         self.mode = mode
@@ -329,9 +329,10 @@ def __getitem__(self, index):
 
         img = 'None'  # Image.open(os.path.join(self.data_dir,'../Image_set_Resize',img_path)).convert("RGB")
 
-        match_label = self.match_label[index]
+        # match_label = self.match_label[index]  # commented: only needed when uncertainty_aware
 
         if self.config['data']['single_uncertainty_aware']:
+            match_label = self.match_label[index]
             img_features = []
             for trial_n in range(self.per_trials):
                 if self.mode == 'train':
@@ -348,6 +349,7 @@ def __getitem__(self, index):
             img_features = torch.stack(img_features)
 
         elif self.config['data']['uncertainty_aware']:
+            match_label = self.match_label[index]
             if self.mode == 'train':
                 if match_label == 0:
                     tag = 'low'
@@ -368,7 +370,8 @@ def __getitem__(self, index):
 
         sample = {
             'idx': index,
-            'eeg': eeg[:, :, self.timesteps[0]:self.timesteps[1]],
+            # 'eeg': eeg[:, :, self.timesteps[0]:self.timesteps[1]],  # original: broken for 2D tensor
+            'eeg': eeg[..., self.timesteps[0]:self.timesteps[1]],  # fixed: use ... to slice last dim
             'label': label,
             'img_path': img_path,
             'img': img,
@@ -377,7 +380,8 @@ def __getitem__(self, index):
             # 'text_features': text_features,
             'session': session,
             'subject': subject,
-            'eeg_mean': eeg_mean[:, :, self.timesteps[0]:self.timesteps[1]],
+            # 'eeg_mean': eeg_mean[:, :, self.timesteps[0]:self.timesteps[1]],  # original: broken for 2D tensor
+            'eeg_mean': eeg_mean[..., self.timesteps[0]:self.timesteps[1]],  # fixed: use ... to slice last dim
         }
         return sample
 

src/retrieval/train.py

@@ -27,6 +27,8 @@
 import matplotlib.pyplot as plt
 
 device = get_device('auto')
+# original: device was int, now torch.device; extract index for Trainer
+device_id = device.index if hasattr(device, 'index') else device
 
 
 def load_model(config, train_loader, test_loader):
@@ -61,7 +63,8 @@ def __init__(self, model, config, train_loader, test_loader, model_type='RN50'):
             self.sim = np.ones(len(train_loader.dataset))
             self.match_label = np.ones(len(train_loader.dataset), dtype=int)
 
-        self.alpha = config['alpha']
+        # self.alpha = config['alpha']  # original: crashes if alpha not in config
+        self.alpha = config.get('alpha', 0.05)  # fixed: default 0.05 when not specified
         self.gamma = 0.3
 
         self.mAP_total = 0
@@ -140,6 +143,8 @@ def forward(self, batch, sample_posterior=False):
 
             return eeg_z, img_z, loss
 
+        else:return eeg_z, img_z, loss
+
     def training_step(self, batch, batch_idx):
         batch_size = batch['idx'].shape[0]
         eeg_z, img_z, loss = self(batch, sample_posterior=True)
@@ -331,9 +336,10 @@ def main(config, yaml):
             mode='min'
         )
 
-    trainer = Trainer(log_every_n_steps=10, #strategy=DDPStrategy(find_unused_parameters=False),
+    # original: Trainer(log_every_n_steps=10, #strategy=DDPStrategy(find_unused_parameters=False), ...)
+    trainer = Trainer(log_every_n_steps=10, enable_progress_bar=True,  # added: show progress bar
                       callbacks=[early_stop_callback, checkpoint_callback], max_epochs=config['train']['epoch'],
-                      devices=[device], accelerator='cuda', logger=logger)
+                      devices=[device_id], accelerator='cuda', logger=logger)
     print(trainer.logger.log_dir)
 
     ckpt_path = 'last'  # None
@@ -393,7 +399,8 @@ def apply_arg_overrides(config, args):
     if args.info is not None:
         config['info'] = args.info
     elif any(value is not None for value in (args.start_time, args.end_time)):
-        config['info'] = f"-ubp{config['alpha']}-[{start_time},{end_time}]"
+        # config['info'] = f"-ubp{config['alpha']}-[{start_time},{end_time}]"  # original: crashes if alpha not in config
+        config['info'] = f"-ubp{config.get('alpha', 0.05)}-[{start_time},{end_time}]"  # fixed: default 0.05
     return config
 
 

src/retrieval/utils.py

@@ -72,7 +72,7 @@ def get_device(gpu_ids):
         selected_device = selected_gpus[0][0]
         # device = torch.device(f'cuda:{selected_device}')
     elif gpu_ids == "cpu":
-        device = torch.device('cpu')
+        selected_device = 'cpu'
     else:
         gpu_ids = list(map(int, gpu_ids.split(",")))
         selected_device = gpu_ids[0]