From 1c001b12feb31da69583803ebc44955c49faa488 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Fri, 6 Nov 2020 17:04:01 +0100 Subject: [PATCH 01/13] Move CustomANISymmetryFunctions construction to ANISymmetryFunctionsOp --- pytorch/SymmetryFunctions.cpp | 33 ++++++++------------------------- 1 file changed, 8 insertions(+), 25 deletions(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index 6b1d1f8..6d13f74 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -115,21 +115,10 @@ class GradANISymmetryFunction : public torch::autograd::Function& EtaR, - const std::vector& ShfR, - const std::vector& EtaA, - const std::vector& Zeta, - const std::vector& ShfA, - const std::vector& ShfZ, - const std::vector& atomSpecies, + const torch::intrusive_ptr& symFunc, const torch::Tensor& positions, const torch::optional& periodicBoxVectors) { - const auto symFunc = torch::intrusive_ptr::make( - numSpecies, Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, atomSpecies, positions); ctx->saved_data["symFunc"] = symFunc; return symFunc->forward(positions, periodicBoxVectors); @@ -141,18 +130,9 @@ class GradANISymmetryFunction : public torch::autograd::Functionbackward(grads); ctx->saved_data.erase("symFunc"); - return { torch::Tensor(), // numSpecies - torch::Tensor(), // Rcr - torch::Tensor(), // Rca - torch::Tensor(), // EtaR - torch::Tensor(), // ShfR - torch::Tensor(), // EtaA - torch::Tensor(), // Zeta - torch::Tensor(), // ShfA - torch::Tensor(), // ShfZ - torch::Tensor(), // atomSpecies - positionsGrad, // positions - torch::Tensor()}; // periodicBoxVectors + return { torch::Tensor(), // symFunc + positionsGrad, // positions + torch::Tensor() }; // periodicBoxVectors }; }; @@ -169,7 +149,10 @@ static torch::autograd::tensor_list ANISymmetryFunctionsOp(int64_t numSpecies, const torch::Tensor& positions, const torch::optional& periodicBoxVectors) { - return GradANISymmetryFunction::apply(numSpecies, Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, atomSpecies, positions, periodicBoxVectors); + const auto symFunc = torch::intrusive_ptr::make( + numSpecies, Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, atomSpecies, positions); + + return GradANISymmetryFunction::apply(symFunc, positions, periodicBoxVectors); } TORCH_LIBRARY(NNPOps, m) { From c8d3d83fa8ed3983c81fbed7163f5404bbc16b90 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Fri, 6 Nov 2020 18:07:59 +0100 Subject: [PATCH 02/13] Move CustomANISymmetryFunctions construction to TorchANISymmetryFunctions.forward --- pytorch/SymmetryFunctions.cpp | 32 ++++++++++++-------------------- pytorch/SymmetryFunctions.py | 11 ++++++++--- 2 files changed, 20 insertions(+), 23 deletions(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index 6d13f74..63a81bf 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -114,17 +114,20 @@ class CustomANISymmetryFunctions : public torch::CustomClassHolder { class GradANISymmetryFunction : public torch::autograd::Function { public: - static torch::autograd::tensor_list forward(torch::autograd::AutogradContext *ctx, - const torch::intrusive_ptr& symFunc, - const torch::Tensor& positions, - const torch::optional& periodicBoxVectors) { + static torch::autograd::tensor_list forward( + torch::autograd::AutogradContext *ctx, + const torch::intrusive_ptr& symFunc, + const torch::Tensor& positions, + const torch::optional& periodicBoxVectors) { ctx->saved_data["symFunc"] = symFunc; return symFunc->forward(positions, periodicBoxVectors); }; - static torch::autograd::tensor_list backward(torch::autograd::AutogradContext *ctx, const torch::autograd::tensor_list& grads) { + static torch::autograd::tensor_list backward( + torch::autograd::AutogradContext *ctx, + const torch::autograd::tensor_list& grads) { const auto symFunc = ctx->saved_data["symFunc"].toCustomClass(); torch::Tensor positionsGrad = symFunc->backward(grads); @@ -136,21 +139,10 @@ class GradANISymmetryFunction : public torch::autograd::Function& EtaR, - const std::vector& ShfR, - const std::vector& EtaA, - const std::vector& Zeta, - const std::vector& ShfA, - const std::vector& ShfZ, - const std::vector& atomSpecies, - const torch::Tensor& positions, - const torch::optional& periodicBoxVectors) { - - const auto symFunc = torch::intrusive_ptr::make( - numSpecies, Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, atomSpecies, positions); +static torch::autograd::tensor_list ANISymmetryFunctionsOp( + const torch::intrusive_ptr& symFunc, + const torch::Tensor& positions, + const torch::optional& periodicBoxVectors) { return GradANISymmetryFunction::apply(symFunc, positions, periodicBoxVectors); } diff --git a/pytorch/SymmetryFunctions.py b/pytorch/SymmetryFunctions.py index 224b793..b783217 100644 --- a/pytorch/SymmetryFunctions.py +++ b/pytorch/SymmetryFunctions.py @@ -29,6 +29,8 @@ from torchani.aev import SpeciesAEV torch.ops.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so')) +torch.classes.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so')) + class TorchANISymmetryFunctions(torch.nn.Module): """Optimized TorchANI symmetry functions @@ -113,10 +115,13 @@ def forward(self, speciesAndPositions: Tuple[Tensor, Tensor], if pbc_ != [True, True, True]: raise ValueError('Only fully periodic systems are supported, i.e. pbc = [True, True, True]') + SymClass = torch.classes.NNPOps.CustomANISymmetryFunctions + symInst = SymClass(self.numSpecies, self.Rcr, self.Rca, self.EtaR, self.ShfR, + self.EtaA, self.Zeta, self.ShfA, self.ShfZ, + species_, positions[0]) + symFunc = torch.ops.NNPOps.ANISymmetryFunctions - radial, angular = symFunc(self.numSpecies, self.Rcr, self.Rca, self.EtaR, self.ShfR, - self.EtaA, self.Zeta, self.ShfA, self.ShfZ, - species_, positions[0], cell) + radial, angular = symFunc(symInst, positions[0], cell) features = torch.cat((radial, angular), dim=1).unsqueeze(0) return SpeciesAEV(species, features) \ No newline at end of file From 3c59e742ac5ee3c1e9dbf1f0cef5eaed79557db0 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Mon, 9 Nov 2020 11:02:14 +0100 Subject: [PATCH 03/13] Move CustomANISymmetryFunctions construction to TorchANISymmetryFunctions.forward --- pytorch/SymmetryFunctions.cpp | 4 ++-- pytorch/SymmetryFunctions.py | 18 +++++++++++------- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index 63a81bf..076f275 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -140,11 +140,11 @@ class GradANISymmetryFunction : public torch::autograd::Function& symFunc, + const torch::optional>& symFunc, const torch::Tensor& positions, const torch::optional& periodicBoxVectors) { - return GradANISymmetryFunction::apply(symFunc, positions, periodicBoxVectors); + return GradANISymmetryFunction::apply(*symFunc, positions, periodicBoxVectors); } TORCH_LIBRARY(NNPOps, m) { diff --git a/pytorch/SymmetryFunctions.py b/pytorch/SymmetryFunctions.py index b783217..f9df293 100644 --- a/pytorch/SymmetryFunctions.py +++ b/pytorch/SymmetryFunctions.py @@ -63,12 +63,13 @@ class TorchANISymmetryFunctions(torch.nn.Module): >>> print(energy, forces) """ + holder: Optional[torch.classes.NNPOps.CustomANISymmetryFunctions] + def __init__(self, symmFunc: torchani.AEVComputer): """ Arguments: symmFunc: the instance of torchani.AEVComputer (https://aiqm.github.io/torchani/api.html#torchani.AEVComputer) """ - super().__init__() self.numSpecies = symmFunc.num_species @@ -81,6 +82,8 @@ def __init__(self, symmFunc: torchani.AEVComputer): self.ShfA = symmFunc.ShfA[0, 0, :, 0].tolist() self.ShfZ = symmFunc.ShfZ[0, 0, 0, :].tolist() + self.holder = None + self.triu_index = torch.tensor([0]) # A dummy variable to make TorchScript happy ;) def forward(self, speciesAndPositions: Tuple[Tensor, Tensor], @@ -102,7 +105,6 @@ def forward(self, speciesAndPositions: Tuple[Tensor, Tensor], species, positions = speciesAndPositions if species.shape[0] != 1: raise ValueError('Batched molecule computation is not supported') - species_: List[int] = species[0].tolist() # Explicit type casting for TorchScript if species.shape + (3,) != positions.shape: raise ValueError('Inconsistent shapes of "species" and "positions"') if cell is not None: @@ -115,13 +117,15 @@ def forward(self, speciesAndPositions: Tuple[Tensor, Tensor], if pbc_ != [True, True, True]: raise ValueError('Only fully periodic systems are supported, i.e. pbc = [True, True, True]') - SymClass = torch.classes.NNPOps.CustomANISymmetryFunctions - symInst = SymClass(self.numSpecies, self.Rcr, self.Rca, self.EtaR, self.ShfR, - self.EtaA, self.Zeta, self.ShfA, self.ShfZ, - species_, positions[0]) + if self.holder is None: + SymClass = torch.classes.NNPOps.CustomANISymmetryFunctions + species_: List[int] = species[0].tolist() # Explicit type casting for TorchScript + self.holder = SymClass(self.numSpecies, self.Rcr, self.Rca, self.EtaR, self.ShfR, + self.EtaA, self.Zeta, self.ShfA, self.ShfZ, + species_, positions) symFunc = torch.ops.NNPOps.ANISymmetryFunctions - radial, angular = symFunc(symInst, positions[0], cell) + radial, angular = symFunc(self.holder, positions[0], cell) features = torch.cat((radial, angular), dim=1).unsqueeze(0) return SpeciesAEV(species, features) \ No newline at end of file From 2831b7773d03efebf775073f426cef35086df78e Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Thu, 12 Nov 2020 17:00:04 +0100 Subject: [PATCH 04/13] Create NNPOps::ANISymmetryFunctions namespace --- pytorch/SymmetryFunctions.cpp | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index 076f275..679b6f0 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -32,6 +32,9 @@ throw std::runtime_error(std::string("Encountered error ")+cudaGetErrorName(result)+" at "+__FILE__+":"+std::to_string(__LINE__));\ } +namespace NNPOps { +namespace ANISymmetryFunctions { + class CustomANISymmetryFunctions : public torch::CustomClassHolder { public: CustomANISymmetryFunctions(int64_t numSpecies_, @@ -105,7 +108,7 @@ class CustomANISymmetryFunctions : public torch::CustomClassHolder { private: torch::TensorOptions tensorOptions; - std::shared_ptr symFunc; + std::shared_ptr<::ANISymmetryFunctions> symFunc; torch::Tensor radial; torch::Tensor angular; torch::Tensor positionsGrad; @@ -163,4 +166,7 @@ TORCH_LIBRARY(NNPOps, m) { .def("forward", &CustomANISymmetryFunctions::forward) .def("backward", &CustomANISymmetryFunctions::backward); m.def("ANISymmetryFunctions", ANISymmetryFunctionsOp); -} \ No newline at end of file +} + +} // namespace ANISymmetryFunctions +} // namespace NNPOps \ No newline at end of file From fd1b3b2ecb1fe75224c6bed20acd5995d89dc274 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Thu, 12 Nov 2020 17:42:42 +0100 Subject: [PATCH 05/13] Simplify names --- pytorch/SymmetryFunctions.cpp | 54 +++++++++++++++++------------------ pytorch/SymmetryFunctions.py | 15 +++++----- 2 files changed, 35 insertions(+), 34 deletions(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index 679b6f0..f3020e0 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -35,19 +35,19 @@ namespace NNPOps { namespace ANISymmetryFunctions { -class CustomANISymmetryFunctions : public torch::CustomClassHolder { +class Holder : public torch::CustomClassHolder { public: - CustomANISymmetryFunctions(int64_t numSpecies_, - double Rcr, - double Rca, - const std::vector& EtaR, - const std::vector& ShfR, - const std::vector& EtaA, - const std::vector& Zeta, - const std::vector& ShfA, - const std::vector& ShfZ, - const std::vector& atomSpecies_, - const torch::Tensor& positions) : torch::CustomClassHolder() { + Holder(int64_t numSpecies_, + double Rcr, + double Rca, + const std::vector& EtaR, + const std::vector& ShfR, + const std::vector& EtaA, + const std::vector& Zeta, + const std::vector& ShfA, + const std::vector& ShfZ, + const std::vector& atomSpecies_, + const torch::Tensor& positions) : torch::CustomClassHolder() { tensorOptions = torch::TensorOptions().device(positions.device()); // Data type of float by default int numAtoms = atomSpecies_.size(); @@ -114,27 +114,27 @@ class CustomANISymmetryFunctions : public torch::CustomClassHolder { torch::Tensor positionsGrad; }; -class GradANISymmetryFunction : public torch::autograd::Function { +class AutogradFunctions : public torch::autograd::Function { public: static torch::autograd::tensor_list forward( torch::autograd::AutogradContext *ctx, - const torch::intrusive_ptr& symFunc, + const torch::intrusive_ptr& holder, const torch::Tensor& positions, const torch::optional& periodicBoxVectors) { - ctx->saved_data["symFunc"] = symFunc; + ctx->saved_data["holder"] = holder; - return symFunc->forward(positions, periodicBoxVectors); + return holder->forward(positions, periodicBoxVectors); }; static torch::autograd::tensor_list backward( torch::autograd::AutogradContext *ctx, const torch::autograd::tensor_list& grads) { - const auto symFunc = ctx->saved_data["symFunc"].toCustomClass(); - torch::Tensor positionsGrad = symFunc->backward(grads); - ctx->saved_data.erase("symFunc"); + const auto holder = ctx->saved_data["holder"].toCustomClass(); + torch::Tensor positionsGrad = holder->backward(grads); + ctx->saved_data.erase("holder"); return { torch::Tensor(), // symFunc positionsGrad, // positions @@ -142,16 +142,16 @@ class GradANISymmetryFunction : public torch::autograd::Function>& symFunc, +torch::autograd::tensor_list operation( + const torch::optional>& holder, const torch::Tensor& positions, const torch::optional& periodicBoxVectors) { - return GradANISymmetryFunction::apply(*symFunc, positions, periodicBoxVectors); + return AutogradFunctions::apply(*holder, positions, periodicBoxVectors); } -TORCH_LIBRARY(NNPOps, m) { - m.class_("CustomANISymmetryFunctions") +TORCH_LIBRARY(NNPOpsANISymmetryFunctions, m) { + m.class_("Holder") .def(torch::init&, // ShfZ const std::vector&, // atomSpecies const torch::Tensor&>()) // positions - .def("forward", &CustomANISymmetryFunctions::forward) - .def("backward", &CustomANISymmetryFunctions::backward); - m.def("ANISymmetryFunctions", ANISymmetryFunctionsOp); + .def("forward", &Holder::forward) + .def("backward", &Holder::backward); + m.def("operation", operation); } } // namespace ANISymmetryFunctions diff --git a/pytorch/SymmetryFunctions.py b/pytorch/SymmetryFunctions.py index f9df293..3275507 100644 --- a/pytorch/SymmetryFunctions.py +++ b/pytorch/SymmetryFunctions.py @@ -31,6 +31,8 @@ torch.ops.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so')) torch.classes.load_library(os.path.join(os.path.dirname(__file__), 'libNNPOpsPyTorch.so')) +Holder = torch.classes.NNPOpsANISymmetryFunctions.Holder +operation = torch.ops.NNPOpsANISymmetryFunctions.operation class TorchANISymmetryFunctions(torch.nn.Module): """Optimized TorchANI symmetry functions @@ -63,7 +65,7 @@ class TorchANISymmetryFunctions(torch.nn.Module): >>> print(energy, forces) """ - holder: Optional[torch.classes.NNPOps.CustomANISymmetryFunctions] + holder: Optional[Holder] def __init__(self, symmFunc: torchani.AEVComputer): """ @@ -118,14 +120,13 @@ def forward(self, speciesAndPositions: Tuple[Tensor, Tensor], raise ValueError('Only fully periodic systems are supported, i.e. pbc = [True, True, True]') if self.holder is None: - SymClass = torch.classes.NNPOps.CustomANISymmetryFunctions species_: List[int] = species[0].tolist() # Explicit type casting for TorchScript - self.holder = SymClass(self.numSpecies, self.Rcr, self.Rca, self.EtaR, self.ShfR, - self.EtaA, self.Zeta, self.ShfA, self.ShfZ, - species_, positions) + self.holder = Holder(self.numSpecies, self.Rcr, self.Rca, + self.EtaR, self.ShfR, + self.EtaA, self.Zeta, self.ShfA, self.ShfZ, + species_, positions) - symFunc = torch.ops.NNPOps.ANISymmetryFunctions - radial, angular = symFunc(self.holder, positions[0], cell) + radial, angular = operation(self.holder, positions[0], cell) features = torch.cat((radial, angular), dim=1).unsqueeze(0) return SpeciesAEV(species, features) \ No newline at end of file From 539b3f2c829aafc5ad17a251ba9b2dc2a162f456 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Thu, 12 Nov 2020 18:14:55 +0100 Subject: [PATCH 06/13] Simplify types --- pytorch/SymmetryFunctions.cpp | 98 ++++++++++++++++++----------------- 1 file changed, 51 insertions(+), 47 deletions(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index f3020e0..eccc4f0 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -35,31 +35,39 @@ namespace NNPOps { namespace ANISymmetryFunctions { +class Holder; +using std::vector; +using HolderPrt = torch::intrusive_ptr; +using torch::Tensor; +using torch::optional; +using Context = torch::autograd::AutogradContext; +using torch::autograd::tensor_list; + class Holder : public torch::CustomClassHolder { public: Holder(int64_t numSpecies_, double Rcr, double Rca, - const std::vector& EtaR, - const std::vector& ShfR, - const std::vector& EtaA, - const std::vector& Zeta, - const std::vector& ShfA, - const std::vector& ShfZ, - const std::vector& atomSpecies_, - const torch::Tensor& positions) : torch::CustomClassHolder() { + const vector& EtaR, + const vector& ShfR, + const vector& EtaA, + const vector& Zeta, + const vector& ShfA, + const vector& ShfZ, + const vector& atomSpecies_, + const Tensor& positions) : torch::CustomClassHolder() { tensorOptions = torch::TensorOptions().device(positions.device()); // Data type of float by default int numAtoms = atomSpecies_.size(); int numSpecies = numSpecies_; - const std::vector atomSpecies(atomSpecies_.begin(), atomSpecies_.end()); + const vector atomSpecies(atomSpecies_.begin(), atomSpecies_.end()); - std::vector radialFunctions; + vector radialFunctions; for (const float eta: EtaR) for (const float rs: ShfR) radialFunctions.push_back({eta, rs}); - std::vector angularFunctions; + vector angularFunctions; for (const float eta: EtaA) for (const float zeta: Zeta) for (const float rs: ShfA) @@ -80,11 +88,11 @@ class Holder : public torch::CustomClassHolder { positionsGrad = torch::empty({numAtoms, 3}, tensorOptions); }; - torch::autograd::tensor_list forward(const torch::Tensor& positions_, const torch::optional& periodicBoxVectors_) { + tensor_list forward(const Tensor& positions_, const optional& periodicBoxVectors_) { - const torch::Tensor positions = positions_.to(tensorOptions); + const Tensor positions = positions_.to(tensorOptions); - torch::Tensor periodicBoxVectors; + Tensor periodicBoxVectors; float* periodicBoxVectorsPtr = nullptr; if (periodicBoxVectors_) { periodicBoxVectors = periodicBoxVectors_->to(tensorOptions); @@ -96,10 +104,10 @@ class Holder : public torch::CustomClassHolder { return {radial, angular}; }; - torch::Tensor backward(const torch::autograd::tensor_list& grads) { + Tensor backward(const tensor_list& grads) { - const torch::Tensor radialGrad = grads[0].clone(); - const torch::Tensor angularGrad = grads[1].clone(); + const Tensor radialGrad = grads[0].clone(); + const Tensor angularGrad = grads[1].clone(); symFunc->backprop(radialGrad.data_ptr(), angularGrad.data_ptr(), positionsGrad.data_ptr()); @@ -109,60 +117,56 @@ class Holder : public torch::CustomClassHolder { private: torch::TensorOptions tensorOptions; std::shared_ptr<::ANISymmetryFunctions> symFunc; - torch::Tensor radial; - torch::Tensor angular; - torch::Tensor positionsGrad; + Tensor radial; + Tensor angular; + Tensor positionsGrad; }; class AutogradFunctions : public torch::autograd::Function { public: - static torch::autograd::tensor_list forward( - torch::autograd::AutogradContext *ctx, - const torch::intrusive_ptr& holder, - const torch::Tensor& positions, - const torch::optional& periodicBoxVectors) { + static tensor_list forward(Context *ctx, + const HolderPrt& holder, + const Tensor& positions, + const optional& periodicBoxVectors) { ctx->saved_data["holder"] = holder; return holder->forward(positions, periodicBoxVectors); }; - static torch::autograd::tensor_list backward( - torch::autograd::AutogradContext *ctx, - const torch::autograd::tensor_list& grads) { + static tensor_list backward(Context *ctx, const tensor_list& grads) { const auto holder = ctx->saved_data["holder"].toCustomClass(); - torch::Tensor positionsGrad = holder->backward(grads); + Tensor positionsGrad = holder->backward(grads); ctx->saved_data.erase("holder"); - return { torch::Tensor(), // symFunc - positionsGrad, // positions - torch::Tensor() }; // periodicBoxVectors + return { Tensor(), // symFunc + positionsGrad, // positions + Tensor() }; // periodicBoxVectors }; }; -torch::autograd::tensor_list operation( - const torch::optional>& holder, - const torch::Tensor& positions, - const torch::optional& periodicBoxVectors) { +tensor_list operation(const optional& holder, + const Tensor& positions, + const optional& periodicBoxVectors) { return AutogradFunctions::apply(*holder, positions, periodicBoxVectors); } TORCH_LIBRARY(NNPOpsANISymmetryFunctions, m) { m.class_("Holder") - .def(torch::init&, // EtaR - const std::vector&, // ShfR - const std::vector&, // EtaA - const std::vector&, // Zeta - const std::vector&, // ShfA - const std::vector&, // ShfZ - const std::vector&, // atomSpecies - const torch::Tensor&>()) // positions + .def(torch::init&, // EtaR + const vector&, // ShfR + const vector&, // EtaA + const vector&, // Zeta + const vector&, // ShfA + const vector&, // ShfZ + const vector&, // atomSpecies + const Tensor&>()) // positions .def("forward", &Holder::forward) .def("backward", &Holder::backward); m.def("operation", operation); From 93c68e0abe6076faacbec20e3d2f7879eb79d2ba Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Wed, 31 Mar 2021 12:49:42 +0200 Subject: [PATCH 07/13] Fix typo --- pytorch/SymmetryFunctions.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index eccc4f0..11102b1 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -37,7 +37,7 @@ namespace ANISymmetryFunctions { class Holder; using std::vector; -using HolderPrt = torch::intrusive_ptr; +using HolderPtr = torch::intrusive_ptr; using torch::Tensor; using torch::optional; using Context = torch::autograd::AutogradContext; @@ -126,7 +126,7 @@ class AutogradFunctions : public torch::autograd::Function { public: static tensor_list forward(Context *ctx, - const HolderPrt& holder, + const HolderPtr& holder, const Tensor& positions, const optional& periodicBoxVectors) { @@ -147,7 +147,7 @@ class AutogradFunctions : public torch::autograd::Function { }; }; -tensor_list operation(const optional& holder, +tensor_list operation(const optional& holder, const Tensor& positions, const optional& periodicBoxVectors) { From 14860f02e96f1e0e1c65e53a9e614801b319e1a5 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Wed, 31 Mar 2021 14:08:58 +0200 Subject: [PATCH 08/13] Implement Holder::is_initialized --- pytorch/SymmetryFunctions.cpp | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index 11102b1..89e352b 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -112,7 +112,11 @@ class Holder : public torch::CustomClassHolder { symFunc->backprop(radialGrad.data_ptr(), angularGrad.data_ptr(), positionsGrad.data_ptr()); return positionsGrad; - } + }; + + bool is_initialized() { + return bool(symFunc); + }; private: torch::TensorOptions tensorOptions; @@ -168,7 +172,8 @@ TORCH_LIBRARY(NNPOpsANISymmetryFunctions, m) { const vector&, // atomSpecies const Tensor&>()) // positions .def("forward", &Holder::forward) - .def("backward", &Holder::backward); + .def("backward", &Holder::backward) + .def("is_initialized", &Holder::is_initialized); m.def("operation", operation); } From 0bef57182996d8133647645ad131a284cad9d457 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Wed, 31 Mar 2021 14:10:34 +0200 Subject: [PATCH 09/13] Don't use Optional[Holder] --- pytorch/SymmetryFunctions.cpp | 3 +++ pytorch/SymmetryFunctions.py | 9 +++++---- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index 89e352b..46b1f67 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -57,6 +57,9 @@ class Holder : public torch::CustomClassHolder { const vector& atomSpecies_, const Tensor& positions) : torch::CustomClassHolder() { + if (numSpecies_ == 0) + return; + tensorOptions = torch::TensorOptions().device(positions.device()); // Data type of float by default int numAtoms = atomSpecies_.size(); int numSpecies = numSpecies_; diff --git a/pytorch/SymmetryFunctions.py b/pytorch/SymmetryFunctions.py index 3275507..af3bb84 100644 --- a/pytorch/SymmetryFunctions.py +++ b/pytorch/SymmetryFunctions.py @@ -65,8 +65,6 @@ class TorchANISymmetryFunctions(torch.nn.Module): >>> print(energy, forces) """ - holder: Optional[Holder] - def __init__(self, symmFunc: torchani.AEVComputer): """ Arguments: @@ -84,7 +82,9 @@ def __init__(self, symmFunc: torchani.AEVComputer): self.ShfA = symmFunc.ShfA[0, 0, :, 0].tolist() self.ShfZ = symmFunc.ShfZ[0, 0, 0, :].tolist() - self.holder = None + # Create an uninitialized holder + self.holder = Holder(0, 0, 0, [], [] , [] , [], [] , [], [], Tensor()) + assert not self.holder.is_initialized() self.triu_index = torch.tensor([0]) # A dummy variable to make TorchScript happy ;) @@ -119,12 +119,13 @@ def forward(self, speciesAndPositions: Tuple[Tensor, Tensor], if pbc_ != [True, True, True]: raise ValueError('Only fully periodic systems are supported, i.e. pbc = [True, True, True]') - if self.holder is None: + if not self.holder.is_initialized(): species_: List[int] = species[0].tolist() # Explicit type casting for TorchScript self.holder = Holder(self.numSpecies, self.Rcr, self.Rca, self.EtaR, self.ShfR, self.EtaA, self.Zeta, self.ShfA, self.ShfZ, species_, positions) + assert self.holder.is_initialized() radial, angular = operation(self.holder, positions[0], cell) features = torch.cat((radial, angular), dim=1).unsqueeze(0) From ab140a5ddbfdfa1a6e2eba91f033399a7dac0728 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Wed, 31 Mar 2021 14:59:04 +0200 Subject: [PATCH 10/13] Fix serializaton --- pytorch/SymmetryFunctions.cpp | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index 46b1f67..e8056e3 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -45,6 +45,11 @@ using torch::autograd::tensor_list; class Holder : public torch::CustomClassHolder { public: + + // Constructor for an uninitialized object + // Note: this is need for serialization + Holder() {}; + Holder(int64_t numSpecies_, double Rcr, double Rca, @@ -57,6 +62,8 @@ class Holder : public torch::CustomClassHolder { const vector& atomSpecies_, const Tensor& positions) : torch::CustomClassHolder() { + // Construct an uninitialized object + // Note: this is needed for Python bindings if (numSpecies_ == 0) return; @@ -148,7 +155,7 @@ class AutogradFunctions : public torch::autograd::Function { Tensor positionsGrad = holder->backward(grads); ctx->saved_data.erase("holder"); - return { Tensor(), // symFunc + return { Tensor(), // holder positionsGrad, // positions Tensor() }; // periodicBoxVectors }; @@ -176,7 +183,15 @@ TORCH_LIBRARY(NNPOpsANISymmetryFunctions, m) { const Tensor&>()) // positions .def("forward", &Holder::forward) .def("backward", &Holder::backward) - .def("is_initialized", &Holder::is_initialized); + .def("is_initialized", &Holder::is_initialized) + .def_pickle( + // __getstate__ + // Note: nothing is during serialization + [](const HolderPtr& self) -> int64_t { return 0; }, + // __setstate__ + // Note: a new uninitialized object is create during deserialization + [](int64_t state) -> HolderPtr { return HolderPtr::make(); } + ); m.def("operation", operation); } From 8676d00312c7d42012f3ec479fe19b63ce25adfa Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Wed, 31 Mar 2021 15:19:34 +0200 Subject: [PATCH 11/13] Update the benckmark --- pytorch/BenchmarkTorchANISymmetryFunctions.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pytorch/BenchmarkTorchANISymmetryFunctions.py b/pytorch/BenchmarkTorchANISymmetryFunctions.py index a2d5e7a..6299c73 100644 --- a/pytorch/BenchmarkTorchANISymmetryFunctions.py +++ b/pytorch/BenchmarkTorchANISymmetryFunctions.py @@ -40,7 +40,7 @@ sum_aev.backward() grad = positions.grad.clone() -N = 40000 +N = 100000 start = time.time() for _ in range(N): aev = symmFunc(speciesPositions).aevs @@ -55,7 +55,5 @@ aev_error = torch.max(torch.abs(aev - aev_ref)) grad_error = torch.max(torch.abs(grad - grad_ref)) -print(aev_error) -print(grad_error) assert aev_error < 0.0002 assert grad_error < 0.007 \ No newline at end of file From a1f9114b7d488437cd571838b58068244ce7cf04 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Wed, 31 Mar 2021 15:32:36 +0200 Subject: [PATCH 12/13] Update the build instructions --- pytorch/README.md | 33 +++++++++++++-------------------- pytorch/environment.yml | 12 ++++++++++++ 2 files changed, 25 insertions(+), 20 deletions(-) create mode 100644 pytorch/environment.yml diff --git a/pytorch/README.md b/pytorch/README.md index c13b3a7..d5c63e6 100644 --- a/pytorch/README.md +++ b/pytorch/README.md @@ -44,39 +44,32 @@ print(energy, forces) ### Build & install -- Crate a *Conda* environment +- Get the source code ```bash -$ conda create -n nnpops \ - -c pytorch \ - -c conda-forge \ - cmake \ - git \ - gxx_linux-64 \ - make \ - mdtraj \ - pytest \ - python=3.8 \ - pytorch=1.6 \ - torchani=2.2 -$ conda activate nnpops +$ git clone https://github.com/openmm/NNPOps.git ``` -- Get the source code + +- Crate a *Conda* environment ```bash -$ git clone https://github.com/peastman/NNPOps.git +$ cd NNPOps +$ conda create -f pytorch/environment.yml +$ conda activate nnpops ``` + - Configure, build, and install ```bash $ mkdir build $ cd build -$ cmake ../NNPOps/pytorch \ +$ cmake ../pytorch \ -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \ -DCMAKE_CUDA_HOST_COMPILER=$CXX \ - -DTorch_DIR=$CONDA_PREFIX/lib/python3.8/site-packages/torch/share/cmake/Torch \ + -DTorch_DIR=$CONDA_PREFIX/lib/python3.9/site-packages/torch/share/cmake/Torch \ -DCMAKE_INSTALL_PREFIX=$CONDA_PREFIX $ make install ``` -- Optional: run tests +- Optional: run tests and benchmarks ```bash -$ cd ../NNPOps/pytorch +$ cd ../pytorch $ pytest TestSymmetryFunctions.py +$ python BenchmarkTorchANISymmetryFunctions.py ``` \ No newline at end of file diff --git a/pytorch/environment.yml b/pytorch/environment.yml new file mode 100644 index 0000000..2cb0894 --- /dev/null +++ b/pytorch/environment.yml @@ -0,0 +1,12 @@ +name: nnpops +channels: + - conda-forge +dependencies: + - cmake + - gxx_linux-64 + - make + - mdtraj + - torchani 2.2 + - pytest + - python 3.9 + - pytorch 1.8.0 \ No newline at end of file From 339e99c9cb6a6c60fd25f5a65885f7e877138206 Mon Sep 17 00:00:00 2001 From: Raimondas Galvelis Date: Wed, 31 Mar 2021 15:47:19 +0200 Subject: [PATCH 13/13] Fix the constructor --- pytorch/SymmetryFunctions.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch/SymmetryFunctions.cpp b/pytorch/SymmetryFunctions.cpp index e8056e3..b4284cc 100644 --- a/pytorch/SymmetryFunctions.cpp +++ b/pytorch/SymmetryFunctions.cpp @@ -48,7 +48,7 @@ class Holder : public torch::CustomClassHolder { // Constructor for an uninitialized object // Note: this is need for serialization - Holder() {}; + Holder() : torch::CustomClassHolder() {}; Holder(int64_t numSpecies_, double Rcr,