AnnaManasyan / VICReg

An easy-to-use implementation of VICReg in Pytorch.
MIT License
19 stars 7 forks source link

Does this code need large GPU memory? #4

Open bqdeng opened 2 years ago

bqdeng commented 2 years ago

Hello! During the training of cifar10 dataset, do you encounter that when the batchsize is set to 2048, you can't run on the dual card nvidia3090? Display memory overflow.

So I changed the batch size to 256, which is still a memory overflow.

Finally, I had no choice but to change it to 128 to run.

However, compared with simclr and swav codes, the batch size that can be set under the same device is not so small. I can generally run 2048 or 1024. Is this normal?

My device is nvidia3090, dual card, with 48g of running video memory. The training data set is cifar10

If you can easily answer, I will be very happy!

AnnaManasyan commented 2 years ago

Hi, Unfortunately, I can't tell what is the issue as I don't have that much gpu memory to test it with a such large batch size. However, in the paper it is mentioned that VICReg consumes less memory than SwAV.