search

yuyangw / MolCLR

Implementation of MolCLR: "Molecular Contrastive Learning of Representations via Graph Neural Networks" in PyG.
MIT License
243 stars 60 forks source link

Apex support #6

Closed peizhenbai closed 2 years ago

peizhenbai commented 2 years ago

Hi @yuyangw

Awesome work about molecule pre-training! I see you provided the apex support for mixed-precision training. Since I have not tried the package before, may I ask whether it is obvious to accelerate model training, and how is the performance gap compared to non-apex support?

Many thanks. Peizhen

yuyangw commented 2 years ago

Hi Peizhen,

Thanks for your kind words. The apex package can train and store the NN model in mixed-precision of float16 and float32. This saves the GPU memory and optimizes the matrix operations during training. Thus, apex accelerates the training and enables larger batch size with limited GPU memory. More details can be found in apex official GitHub repo: https://github.com/NVIDIA/apex.

Hope this helps!

Best, Yuyang

peizhenbai commented 2 years ago

Thanks for your quick reply, Yuyang! It's very helpful :-)