NT_Xent Loss function: all negatives are not being used in multi-gpu case ?

sthalles / SimCLR

PyTorch implementation of SimCLR: A Simple Framework for Contrastive Learning of Visual Representations

https://sthalles.github.io/simple-self-supervised-learning/

MIT License

2.19k stars 457 forks source link

NT_Xent Loss function: all negatives are not being used in multi-gpu case ? #59

Open HuangChiEn opened 1 year ago

HuangChiEn commented 1 year ago

@sthalles Hello, i saw there have the similar issue before, the problem is i haven't saw any torch.gather API to collect the negative terms crossing all gpu in loss function calculation. So, is that means this repo does not support the multi-gpu yet ?

Any suggestion will be appreciated !!