diff --git a/configs/baseline.yaml b/configs/baseline.yaml index 6b31971..ee0faa2 100644 --- a/configs/baseline.yaml +++ b/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 diff --git a/configs/baseline_ubp.yaml b/configs/baseline_ubp.yaml index 03bca3b..c659136 100644 --- a/configs/baseline_ubp.yaml +++ b/configs/baseline_ubp.yaml @@ -1,15 +1,14 @@ -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: @@ -17,7 +16,7 @@ models: 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' \ No newline at end of file diff --git a/configs/baseline_ubp_GX.yaml b/configs/baseline_ubp_GX.yaml index 1e6fe91..4a7a2c1 100644 --- a/configs/baseline_ubp_GX.yaml +++ b/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: diff --git a/configs/baseline_ubp_sdxl.yaml b/configs/baseline_ubp_sdxl.yaml index c739c4a..5028bc6 100644 --- a/configs/baseline_ubp_sdxl.yaml +++ b/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: diff --git a/configs/baseline_ubp_sota.yaml b/configs/baseline_ubp_sota.yaml index 12c1350..faadd6e 100644 --- a/configs/baseline_ubp_sota.yaml +++ b/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' + diff --git a/configs/baseline_ubp_vae.yaml b/configs/baseline_ubp_vae.yaml index ce241d0..822047c 100644 --- a/configs/baseline_ubp_vae.yaml +++ b/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: diff --git a/src/retrieval/data_eeg.py b/src/retrieval/data_eeg.py index efc29a9..a1feb70 100644 --- a/src/retrieval/data_eeg.py +++ b/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 diff --git a/src/retrieval/train.py b/src/retrieval/train.py index 1fe06d5..efbd891 100644 --- a/src/retrieval/train.py +++ b/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 diff --git a/src/retrieval/utils.py b/src/retrieval/utils.py index b7a45a2..52a37c2 100644 --- a/src/retrieval/utils.py +++ b/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]