AstraZeneca/chemicalx

GPU Transfer

benedekrozemberczki opened this issue · 3 comments

  • Accelerating training by transfer to GPU.

@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?

@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?

draft PR: #76