Question about the hyperparamter used to train the model

[beotborry]

Hi, first of all, congratulations on your work accepted to ICML!

I have a few questions about the code details to reproduce your reported results.

First, I have found some mismatches between the code and the paper.

1. It is written that \mathcal{L}_{lowlevel} is used with \alpha_2 of 0.016 in your paper. However, I have found that --no_blurry_recon argument is used in the accel.slurm file and the value of blur_scale is 0.5 in the file.
2. It is written that \alpha_1 value is 0.033 but the value of clip_scale is 1 in the accel.slurm file.
3. There is no scalar for \mathcal{L}_{prior} but the value of prior_scale is 30 in the accel.slurm file.
4. There is no scalar for \mathcal{L}_{softclip}(VIC, \hat VIC) but I have found that the loss is multiplied by 0.1

            if blurry_recon:     
                image_enc_pred, transformer_feats = blurry_image_enc_

                image_enc = autoenc.encode(2*image-1).latent_dist.mode() * 0.18215
                loss_blurry = l1(image_enc_pred, image_enc)
                loss_blurry_total += loss_blurry.item()

                if epoch < int(mixup_pct * num_epochs):
                    image_enc_shuf = image_enc[perm]
                    betas_shape = [-1] + [1]*(len(image_enc.shape)-1)
                    image_enc[select] = image_enc[select] * betas[select].reshape(*betas_shape) + \
                        image_enc_shuf[select] * (1 - betas[select]).reshape(*betas_shape)

                image_norm = (image - mean)/std
                image_aug = (blur_augs(image) - mean)/std
                _, cnx_embeds = cnx(image_norm)
                _, cnx_aug_embeds = cnx(image_aug)

                cont_loss = utils.soft_cont_loss(
                    nn.functional.normalize(transformer_feats.reshape(-1, transformer_feats.shape[-1]), dim=-1),
                    nn.functional.normalize(cnx_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),
                    nn.functional.normalize(cnx_aug_embeds.reshape(-1, cnx_embeds.shape[-1]), dim=-1),
                    temp=0.2)
                loss_blurry_cont_total += cont_loss.item()

                loss += (loss_blurry + 0.1*cont_loss) * blur_scale #/.18215

So, I wonder which one should I have to use to reproduce your results.

Secondly, I wonder if you adjusted learning rate or epochs to compensate the reduced iterations when using multiple gpus.
For example, it says that the batch_size is equal to 24 for single subject training in the accel.slurm file, but is it the same for the 4gpus? (in other words, is the global batch size equal to 24*4=96 for 4gpus?)
In addition, the value of max_lr is 3e-4, but does it multiplied by 4 for the 4gpus?

Thank you in advance.
Juhyeon Park

[PaulScotti]

Hi Juhyeon thanks for the interest in our work!

The alphas you are referring to were rewritten in the paper to be relative to the largest scale value. That is, prior_scale = 30 blur_scale = 0.5 and clip_scale = 1.0 are correct but for writing the paper we divided all these numbers by the highest value of 30. If you change the scales to prior_scale = 30/30, blur_scale = .5/30, clip_scale = 1/30 you'll get the numbers as reflected in the manuscript.

There is no scalar for \mathcal{L}_{softclip}(VIC, \hat VIC) but I have found that the loss is multiplied by 0.1

This is correct, we neglected to mention in the manuscript that we hard-coded a scalar to multiply this loss by 0.1.

The code in the github reflects what was used to train the model, so if you are looking to replicate I'd advise to follow the code as opposed to the manuscript equations.

Secondly, I wonder if you adjusted learning rate or epochs to compensate the reduced iterations when using multiple gpus.
For example, it says that the batch_size is equal to 24 for single subject training in the accel.slurm file, but is it the same for the 4gpus? (in other words, is the global batch size equal to 24*4=96 for 4gpus?)
In addition, the value of max_lr is 3e-4, but does it multiplied by 4 for the 4gpus?

The code does adjust the batch size and number of epochs depending on the number of gpus, but it does not adjust learning rate.

batch_size = int(os.environ["GLOBAL_BATCH_SIZE"]) // num_devices

[beotborry]

Thank you for your response.

If possible, could you share the script file you used to train the final model (trained with 8 A100 gpus) for the reported values in the paper?

I'm struggling to reproduce your results, so I want to check training details (lr, batch size etc.) with your script.

Thank you.
Juhyeon Park

[PaulScotti]

What part of reproducing are you finding problems with? We have already had collaborators use our code to successfully reproduce our results, so I believe the code on the GitHub should be in working order. It has only been minimally modified from the actual original code we used for the paper.

For the exact original code you can attempt to browse through my public wandb runs: https://stability.wandb.io/paul-scotti/mindeye?nw=nwuserpaulscotti

[beotborry]

I'm facing a problem that the performance degrades when I train a model with multiple gpus.

I have setup the configurations of accelerator as you did.

I will check your shared wandb log and get back to you.

I appreciate your help!

[PaulScotti]

OK yes please let me know if you find a problem where multi-gpu is leading to performance problems!

[beotborry]

Well, I cloned your repository and tried to train a model for subject 1, w/o pretrained model, with 1 session of data.
The training works with 1 gpu setting, but the process of training stucks when multiple gpus are used.
Could you check this problem?
Thank you.

----- Update -----
I resolved this issue by increasing num_workers of train dataloader.
I will try again and get back to you!

[beotborry]

@PaulScotti
I trained a model with 1 session of subj1 data without pretrained model.

final_subj01_notpretrained_1sess_24bs is trained with a single gpu, final_subj01_notpretrained_1sees_24bs_2gpu is trained with 2 gpus and batch size of 24 (i.e. global batch size is equal to 48) and
final_subj01_notpretrained_1sees_12bs_2gpu is trained with 2gpus and batch size of 12 (i.e. global batch size is equal to 24)

I used learning rate of 3e-4 for training the models.

The test log is as follows.

It seems that all of test metrics degrades when multiple gpus are used.

[PaulScotti]

I see I was thinking you were referring to slight performance degradation with multi-gpu but these are huge differences!

I think you must have some sort of issue in your multi-gpu setup in that case, you can take a look as some of my public wandb runs and see that running the model with 1 sess and 8 gpus trains fine -- https://stability.wandb.io/paul-scotti/mindeye/runs/subj01_pretrained_1sess_1024hid_nolow_seed8

I'd advise you do some debugging with your setup to ensure every batch is indeed being handled properly by each gpu in the multi-gpu case, also look into ensuring your accelerate config is properly configured as discussed in the readme

[HaoyyLi]

I also found the above problem when using multi-gpu training, regardless of whether I changed the learning rate and global batch size. May I ask if you have solved this problem now?

[beotborry]

Yes, I've solved the problem now. In my case, the critical factor was the version of accelerator.
I'm not sure, but how the learning rate is scheduled seems quite different across the accelerator version.
When I install the accelerate==0.24.1, the problem is solved.

[zjsndy]

"Hi, may I ask if you’ve ever encountered NaN issues when using mixed-precision training? For me, the only way to fix it so far is to switch to float32."