A better PyTorch wrapper

[raimis]

• Refactor the wrapper to reuse ANISymmetryFunctions object
• Fix serialization
• Update the benchmarks
• Update the build instructions

[peastman]

Maybe call it something more descriptive like symmetryFunctions? In my example, the name holder was a ValueHolder object.

[raimis]

I'm still working on cleaning up the code, but to call an instance symmetryFunctions is misleading and might be confused with an actual reference to the function symFunc.

[raimis]

Benchmarks

Molecule: 46 atoms (pytorch/molecules/2iuz_ligand.mol2)
Model: ANI-2x
GPU: GTX 1080 Ti
Script: pytorch/BenchmarkTorchANISymmetryFunctions.py
Mode: forward and backward passes

Original TorchANI: 5.0 ms
Current NNPOps: 1.1 ms
New NNPOps: 0.4 ms

[peastman]

Very nice!

[peastman]

What's the status of this PR? It would be really good if we could get it finished up and merged in.

[raimis]

It works, but needs some tests. I will finish it after more or less two weeks, when the paper is done.

[peastman]

Any more updates? Until this is done, we can't use this kernel in OpenMM.

[raimis]

Ping @giadefa

[raimis]

@peastman I have managed to fix the serialization issue. Could you review?

[peastman]

Looks much better! Can we use a more meaningful name than "Holder", for example, ANISymmetryFunctions? That will make the code easier to understand.

Have you tested this to see if compiling to TorchScript works?

[raimis]

Looks much better! Can we use a more meaningful name than "Holder", for example, ANISymmetryFunctions? That will make the code easier to understand.

PyTorch wrapper implement four classes/functions:

• NNPOps::ANISymmetryFunctions::Holder (Python side: torch.classes.NNPOpsANISymmetryFunctions.Holder)
• NNPOps::ANISymmetryFunctions::AutogradFunctions::forward
• NNPOps::ANISymmetryFunctions::AutogradFunctions::backward
• NNPOps::ANISymmetryFunctions::operation (Python side: torch.ops.NNPOpsANISymmetryFunctions.operation)

I doubt that renaming NNPOps::ANISymmetryFunctions::Holder to NNPOps::ANISymmetryFunctions::ANISymmetryFunctions will make code easier to understand. Anyone wanting understand what is going on here most likely will have to refer to PyTorch C++ API documentation, which uses such class names as "Holder". Also, there is already ANISymmetryFunctions class in ani/ANISymmetryFunctions.h

Have you tested this to see if compiling to TorchScript works?

Yes, test_model_serialization in pytorch/TestSymmetryFunctions.py does that.

[giadefa]

can we merge this?

[peastman]

Ok!