Subclass ModuleList to simplify code (#385)

aiqm · Nov 13, 2019 · e41c2d9 · e41c2d9
1 parent 9337213
commit e41c2d9
Showing 1 changed file with 10 additions and 19 deletions.
diff --git a/torchani/nn.py b/torchani/nn.py
@@ -8,7 +8,7 @@ class SpeciesEnergies(NamedTuple):
     energies: Tensor
 
 
-class ANIModel(torch.nn.Module):
+class ANIModel(torch.nn.ModuleList):
     """ANI model that compute energies from species and AEVs.
 
     Different atom types might have different modules, when computing
@@ -27,11 +27,7 @@ class ANIModel(torch.nn.Module):
     """
 
     def __init__(self, modules):
-        super(ANIModel, self).__init__()
-        self.module_list = torch.nn.ModuleList(modules)
-
-    def __getitem__(self, i):
-        return self.module_list[i]
+        super(ANIModel, self).__init__(modules)
 
     def forward(self, species_aev: Tuple[Tensor, Tensor],
                 cell: Optional[Tensor] = None,
@@ -44,7 +40,7 @@ def forward(self, species_aev: Tuple[Tensor, Tensor],
 
         output = aev.new_zeros(species_.shape)
 
-        for i, m in enumerate(self.module_list):
+        for i, m in enumerate(self):
             mask = (species_ == i)
             midx = mask.nonzero().flatten()
             if midx.shape[0] > 0:
@@ -54,40 +50,35 @@ def forward(self, species_aev: Tuple[Tensor, Tensor],
         return SpeciesEnergies(species, torch.sum(output, dim=1))
 
 
-class Ensemble(torch.nn.Module):
+class Ensemble(torch.nn.ModuleList):
     """Compute the average output of an ensemble of modules."""
 
     def __init__(self, modules):
-        super(Ensemble, self).__init__()
-        self.modules_list = torch.nn.ModuleList(modules)
-        self.size = len(self.modules_list)
+        super(Ensemble, self).__init__(modules)
+        self.size = len(modules)
 
     def forward(self, species_input: Tuple[Tensor, Tensor],
                 cell: Optional[Tensor] = None,
                 pbc: Optional[Tensor] = None) -> SpeciesEnergies:
         assert cell is None
         assert pbc is None
         sum_ = 0
-        for x in self.modules_list:
+        for x in self:
             sum_ += x(species_input)[1]
         species, _ = species_input
         return SpeciesEnergies(species, sum_ / self.size)
 
-    def __getitem__(self, i):
-        return self.modules_list[i]
-
 
-class Sequential(torch.nn.Module):
+class Sequential(torch.nn.ModuleList):
     """Modified Sequential module that accept Tuple type as input"""
 
     def __init__(self, *modules):
-        super(Sequential, self).__init__()
-        self.modules_list = torch.nn.ModuleList(modules)
+        super(Sequential, self).__init__(modules)
 
     def forward(self, input_: Tuple[Tensor, Tensor],
                 cell: Optional[Tensor] = None,
                 pbc: Optional[Tensor] = None):
-        for module in self.modules_list:
+        for module in self:
             input_ = module(input_, cell=cell, pbc=pbc)
             cell = None
             pbc = None