aiqm / torchani

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

CUAEV refactor to initialize aev parameter only once #575

Closed yueyericardo closed 3 years ago

yueyericardo commented 3 years ago

original result:

RUN                        Total AEV    Forward      Backward     Force        Optimizer    Others       Epoch time   GPU
0 cu Energy train          21.9 ms      13.1 ms      15.5 ms      0.0 ms       11.0 ms      3.6 ms       65.2 ms      1315.5MB
1 py Energy train          80.7 ms      12.9 ms      16.5 ms      0.0 ms       13.0 ms      3.7 ms       126.8 ms     2123.5MB
2 cu Energy + Force train  21.7 ms      12.9 ms      131.8 ms     75.2 ms      6.7 ms       5.0 ms       253.3 ms     2761.5MB
3 py Energy + Force train  80.2 ms      12.9 ms      435.0 ms     167.4 ms     6.9 ms       5.4 ms       707.8 ms     5667.5MB

original result (only forward):

File: small.pdb, Molecule size: 264

Original TorchANI:
  Duration: 0.42 s
  Speed: 2.11 ms/it

CUaev:
  Duration: 0.08 s
  Speed: 0.40 ms/it
  Speed up: 5.24 X

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

File: 1hz5.pdb, Molecule size: 973

/home/richard/dev/torchani/torchani/resources/
Original TorchANI:
  Duration: 0.54 s
  Speed: 2.70 ms/it

CUaev:
  Duration: 0.19 s
  Speed: 0.93 ms/it
  Speed up: 2.92 X

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

File: 6W8H.pdb, Molecule size: 3410

/home/richard/dev/torchani/torchani/resources/
Original TorchANI:
  Duration: 1.37 s
  Speed: 6.83 ms/it

CUaev:
  Duration: 0.94 s
  Speed: 4.68 ms/it
  Speed up: 1.46 X

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

original result (forward, backward, doublebackward):

File: small.pdb, Molecule size: 264

Original TorchANI:
  Duration: 3.19 s
  Speed: 15.95 ms/it

CUaev:
  Duration: 0.49 s
  Speed: 2.46 ms/it
  Speed up: 6.49 X

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

File: 1hz5.pdb, Molecule size: 973

/home/richard/dev/torchani/torchani/resources/
Original TorchANI:
  Duration: 3.63 s
  Speed: 18.14 ms/it

CUaev:
  Duration: 0.70 s
  Speed: 3.49 ms/it
  Speed up: 5.19 X

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

File: 6W8H.pdb, Molecule size: 3410

/home/richard/dev/torchani/torchani/resources/
Original TorchANI:
  Duration: 4.62 s
  Speed: 23.10 ms/it

CUaev:
  Duration: 1.46 s
  Speed: 7.32 ms/it
  Speed up: 3.16 X

----------------------------------------------------------------------
lgtm-com[bot] commented 3 years ago

This pull request introduces 1 alert when merging 20a3eec597f60f6ee3bf227954a3e1836f479422 into 191cf2d1dd72b85d7ac7fa6ca7c719e5068a0a3d - view on LGTM.com

new alerts:

yueyericardo commented 3 years ago

Some notes about the memory issue are saved at: Memory blow up when return a tensor from torch::CustomClassHolder by yueyericardo · Pull Request #1 · zasdfgbnm/test-autograd