AstraZeneca / chemicalx

A PyTorch and TorchDrug based deep learning library for drug pair scoring. (KDD 2022)
https://chemicalx.readthedocs.io
Apache License 2.0
715 stars 87 forks source link

GPU Transfer #65

Closed benedekrozemberczki closed 2 years ago

benedekrozemberczki commented 2 years ago
GavEdwards commented 2 years ago

@mughetto @benedekrozemberczki A simple option for flexible GPU support is to use Accelerate.

image

I played around with a quick implementation for chemicalx and it's very easy, only changed a few lines of code. I tested on CPU and single GPU and it works well so far. GPU speeds up training and gets high utilisation, e.g the deepddi example is so much faster with GPU enabled 🏎.

If you're interested in this approach I can make an initial PR?

benedekrozemberczki commented 2 years ago

@GavEdwards Sounds great! Right know, we really struggle with speed for the experiments. Can you make sure that the molecule and fingerprint based models both work on GPU?

GavEdwards commented 2 years ago

draft PR: https://github.com/AstraZeneca/chemicalx/pull/76