openmm / NNPOps

High-performance operations for neural network potentials
Other
79 stars 17 forks source link

GPU memory consumption of NNPOps #77

Closed wiederm closed 1 year ago

wiederm commented 1 year ago

I have seen increased GPU memory consumption using NNPOps compared to the original torchani implementation. Attached is a plot that shows the memory consumption of ANI-2x on the CUDA platform with nnpops and the torchani implementation for a watercluster test system. I am wondering if this is expected behavior or if I am doing something wrong? I am using the openmm-8.0-beta conda environment and openmmml to perform the benchmark watercluster calculations.

(The plot shows on the x-axis the number of atoms in the system and on the left y-axis the time it takes to execute a single energy evaluation and on the right y-axis the GPU memory footprint of the model. The dotted lines show the GPU memory. The labels are always model/platform/implementation)

benchmarking

raimis commented 1 year ago

Yes, it is expected. The nnpops implementation is trading memory for speed.

wiederm commented 1 year ago

Thanks for the clarification!