Zasder3 / train-CLIP

A PyTorch Lightning solution to training OpenAI's CLIP from scratch.
MIT License
653 stars 78 forks source link

Multi GPU training #7

Closed niatzt closed 3 years ago

niatzt commented 3 years ago

Thanks for sharing the code.

I am not familiar with Lightning. It seems that the Code supports multiGPU (https://github.com/Zasder3/train-CLIP/blob/8d454de1999af4be93b6d99b956e83f005ff70dd/models/wrapper.py#L43), but I am not sure how to initiate the multi-GPU training.

Besides, just to confirm, the code does not initialize the weights using the pretrained model?

Zasder3 commented 3 years ago

To use multiple GPUs (as long as they’re on the same device) all you need to do is pass the argument —gpus N into the command line argument when you run the script. N being the number of GPUs on the system. If you are trying to use multiple clusters I’d recommend reading the documentation.

You are correct this is not a pretrained initialization! These are random weights. I’d love to know how this works out for you as I haven’t been able to properly test any of the Multi-GPU code.

tchaton commented 3 years ago

Hey @niatzt,

Yes, PyTorch Lightning makes it really simple to work with multi-gpu in lot of different configurations.

Here is the doc for Multi GPUS: https://pytorch-lightning.readthedocs.io/en/stable/advanced/multi_gpu.html

Also, we support DeepSpeed, DDP Sharded, Horovod and many more :)

Best, T.C

tchaton commented 3 years ago

Dear @Zasder3,

First of all, this repo looks really neat ! Awesome work !

Would you be interested in upstreaming your work to Lightning Flash directly: https://github.com/PyTorchLightning/lightning-flash ?

Here is our Template Task to get you started: https://lightning-flash.readthedocs.io/en/latest/template/intro.html.

This would be our first multi modal task and you can work the Lightning Team to make the experience super neat :)

Best, T.C

Zasder3 commented 3 years ago

Hi @tchaton,

I'd love to contribute the code! I'm currently in the process of adding one more large feature to the repo, but then I'd love to add it to flash.

Zasder3 commented 3 years ago

@niatzt Multi-GPU training is in working order now! Just specify --gpus [however many you want to use] --accelerator ddp and you should be in business!