aiqm / torchani

Accurate Neural Network Potential on PyTorch
https://aiqm.github.io/torchani/
MIT License
459 stars 126 forks source link

Fix CUAEV slower than pyaev as molecule size increase to 1000 atoms or more #565

Closed yueyericardo closed 3 years ago

yueyericardo commented 3 years ago

Related issue: #522

yueyericardo commented 3 years ago

Benchmark on TITAN V

Before

File: small.pdb, Molecule size: 264
Original TorchANI:
  Duration: 1.26 s
  Speed: 2.53 ms/it

CUaev:
  Duration: 0.25 s
  Speed: 0.51 ms/it
  Speed up: 4.97 X

----------------------------------------------------------------------

File: 1hz5.pdb, Molecule size: 973

Original TorchANI:
  Duration: 1.21 s
  Speed: 2.43 ms/it

CUaev:
  Duration: 2.42 s
  Speed: 4.83 ms/it
  Speed up (slower): 0.50 X

----------------------------------------------------------------------

File: 6W8H.pdb, Molecule size: 3410

Original TorchANI:
  Duration: 1.86 s
  Speed: 3.72 ms/it

CUaev:
  Duration: 27.69 s
  Speed: 55.38 ms/it
  Speed up (slower): 0.07 X

After

File: small.pdb, Molecule size: 264
  Duration: 1.37 s
  Speed: 2.74 ms/it

CUaev:
  Duration: 0.20 s
  Speed: 0.39 ms/it
  Speed up: 7.03 X

----------------------------------------------------------------------

File: 1hz5.pdb, Molecule size: 973

Original TorchANI:
  Duration: 1.34 s
  Speed: 2.69 ms/it

CUaev:
  Duration: 0.28 s
  Speed: 0.56 ms/it
  Speed up: 4.82 X

----------------------------------------------------------------------

File: 6W8H.pdb, Molecule size: 3410

Original TorchANI:
  Duration: 2.33 s
  Speed: 4.67 ms/it

CUaev:
  Duration: 1.44 s
  Speed: 2.88 ms/it
  Speed up: 1.62 X

----------------------------------------------------------------------