Delete nn.py and create torchani.models to support rich types of models (#26)

tests/test_benchmark.py

@@ -93,11 +93,11 @@ class TestBenchmark(unittest.TestCase):
     def testModelOnAEV(self):
         aev_computer = torchani.SortedAEV(
             dtype=self.dtype, device=self.device)
-        model = torchani.ModelOnAEV(
-            aev_computer, benchmark=True, from_nc=None)
+        model = torchani.models.NeuroChemNNP(
+            aev_computer, benchmark=True)
         self._testModule(model, ['forward>aev', 'forward>nn'])
-        model = torchani.ModelOnAEV(
-            aev_computer, benchmark=True, derivative=True, from_nc=None)
+        model = torchani.models.NeuroChemNNP(
+            aev_computer, benchmark=True, derivative=True)
         self._testModule(
             model, ['forward>aev', 'forward>nn', 'forward>derivative'])
 

tests/test_energies.py

@@ -16,7 +16,7 @@ class TestEnergies(unittest.TestCase):
         self.tolerance = 5e-5
         self.aev_computer = torchani.SortedAEV(
             dtype=dtype, device=torch.device('cpu'))
-        self.nnp = torchani.ModelOnAEV(self.aev_computer, from_nc=None)
+        self.nnp = torchani.models.NeuroChemNNP(self.aev_computer)
 
     def _test_molecule(self, coordinates, species, energies):
         shift_energy = torchani.EnergyShifter(torchani.buildin_sae_file)

tests/test_ensemble.py

@@ -15,15 +15,17 @@ class TestEnsemble(unittest.TestCase):
         self.conformations = 20
 
     def _test_molecule(self, coordinates, species):
+        n = torchani.buildin_ensemble
         prefix = torchani.buildin_model_prefix
-        n = torchani.buildin_ensembles
         aev = torchani.SortedAEV(device=torch.device('cpu'))
         coordinates, species = aev.sort_by_species(coordinates, species)
-        ensemble = torchani.ModelOnAEV(aev, derivative=True,
-                                       from_nc=prefix,
-                                       ensemble=n)
-        models = [torchani.ModelOnAEV(aev, derivative=True,
-                  from_nc=prefix + '{}/networks/'.format(i)) for i in range(n)]
+        ensemble = torchani.models.NeuroChemNNP(aev, derivative=True,
+                                                ensemble=True)
+        models = [torchani.models.
+                  NeuroChemNNP(aev, derivative=True,
+                               ensemble=False,
+                               from_=prefix + '{}/networks/'.format(i))
+                  for i in range(n)]
 
         energy1, force1 = ensemble(coordinates, species)
         energy2, force2 = zip(*[m(coordinates, species) for m in models])

tests/test_forces.py

@@ -15,8 +15,8 @@ class TestForce(unittest.TestCase):
         self.tolerance = 1e-5
         self.aev_computer = torchani.SortedAEV(
             dtype=dtype, device=torch.device('cpu'))
-        self.nnp = torchani.ModelOnAEV(
-            self.aev_computer, derivative=True, from_nc=None)
+        self.nnp = torchani.models.NeuroChemNNP(
+            self.aev_computer, derivative=True)
 
     def _test_molecule(self, coordinates, species, forces):
         _, derivative = self.nnp(coordinates, species)

torchani/__init__.py

 from .energyshifter import EnergyShifter
-from .nn import ModelOnAEV, PerSpeciesFromNeuroChem
+from . import models
 from .aev import SortedAEV
 from .env import buildin_const_file, buildin_sae_file, buildin_network_dir, \
-    buildin_model_prefix, buildin_ensembles, default_dtype, default_device
+    buildin_model_prefix, buildin_ensemble, default_dtype, default_device
 
-__all__ = ['SortedAEV', 'EnergyShifter', 'ModelOnAEV',
-           'PerSpeciesFromNeuroChem', 'data', 'buildin_const_file',
-           'buildin_sae_file', 'buildin_network_dir', 'buildin_model_prefix',
-           'buildin_ensembles', 'default_dtype', 'default_device']
+__all__ = ['SortedAEV', 'EnergyShifter', 'models', 'data',
+           'buildin_const_file', 'buildin_sae_file', 'buildin_network_dir',
+           'buildin_model_prefix', 'buildin_ensemble',
+           'default_dtype', 'default_device']

torchani/_six.py

+# file for python 2 compatibility
+
+import math
+
+if not hasattr(math, 'inf'):
+    math.inf = float('inf')

torchani/env.py

@@ -14,7 +14,7 @@ buildin_network_dir = pkg_resources.resource_filename(
 buildin_model_prefix = pkg_resources.resource_filename(
     __name__, 'resources/ani-1x_dft_x8ens/train')
 
-buildin_ensembles = 8
+buildin_ensemble = 8
 
 default_dtype = torch.float32
 default_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

torchani/models/__init__.py

+from .custom import CustomModel
+from .neurochem_nnp import NeuroChemNNP
+
+__all__ = ['CustomModel', 'NeuroChemNNP']

torchani/models/ani_model.py

+from ..aev_base import AEVComputer
+import torch
+from ..benchmarked import BenchmarkedModule
+
+
+class ANIModel(BenchmarkedModule):
+    """Subclass of `torch.nn.Module` for the [xyz]->[aev]->[per_atom_y]->y
+    pipeline.
+
+    Attributes
+    ----------
+    aev_computer : AEVComputer
+        The AEV computer.
+    output_length : int
+        The length of output vector
+    suffixes : sequence
+        Different suffixes denote different models in an ensemble.
+    model_<X><suffix> : nn.Module
+        Model of suffix <suffix> for species <X>. There should be one such
+        attribute for each supported species.
+    reducer : function
+        Function of (input, dim)->output that reduce the input tensor along the
+        given dimension to get an output tensor. This function will be called
+        with the per atom output tensor with internal shape as input, and
+        desired reduction dimension as dim, and should reduce the input into
+        the tensor containing desired output.
+    output_length : int
+        Length of output of each submodel.
+    derivative : boolean
+        Whether to support computing the derivative w.r.t coordinates,
+        i.e. d(output)/dR
+    derivative_graph : boolean
+        Whether to generate a graph for the derivative. This would be required
+        only if the derivative is included as part of the loss function.
+    timers : dict
+        Dictionary storing the the benchmark result. It has the following keys:
+            aev : time spent on computing AEV.
+            nn : time spent on computing output from AEV.
+            derivative : time spend on computing derivative w.r.t. coordinates
+                after the outputs is given. This key is only available if
+                derivative computation is turned on.
+            forward : total time for the forward pass
+    """
+
+    def __init__(self, aev_computer, suffixes, reducer, output_length, models,
+                 derivative=False, derivative_graph=False, benchmark=False):
+        super(ANIModel, self).__init__(benchmark)
+        if not isinstance(aev_computer, AEVComputer):
+            raise TypeError(
+                "ModelOnAEV: aev_computer must be a subclass of AEVComputer")
+        self.aev_computer = aev_computer
+
+        self.suffixes = suffixes
+        self.reducer = reducer
+        self.output_length = output_length
+        for i in models:
+            setattr(self, i, models[i])
+
+        self.derivative = derivative
+        if not derivative and derivative_graph:
+            raise ValueError(
+                '''BySpeciesModel: can not create graph for derivative if the
+                computation of derivative is turned off''')
+        self.derivative_graph = derivative_graph
+        if derivative and self.output_length != 1:
+            raise ValueError(
+                'derivative can only be computed for output length 1')
+
+        if benchmark:
+            self.compute_aev = self._enable_benchmark(self.compute_aev, 'aev')
+            self.aev_to_output = self._enable_benchmark(
+                self.aev_to_output, 'nn')
+            if derivative:
+                self.compute_derivative = self._enable_benchmark(
+                    self.compute_derivative, 'derivative')
+            self.forward = self._enable_benchmark(self.forward, 'forward')
+
+    def compute_aev(self, coordinates, species):
+        """Compute full AEV
+
+        Parameters
+        ----------
+        coordinates : torch.Tensor
+            The pytorch tensor of shape (conformations, atoms, 3) storing
+            the coordinates of all atoms of all conformations.
+        species : list of string
+            List of string storing the species for each atom.
+
+        Returns
+        -------
+        torch.Tensor
+            Pytorch tensor of shape (conformations, atoms, aev_length) storing
+            the computed AEVs.
+        """
+        radial_aev, angular_aev = self.aev_computer(coordinates, species)
+        fullaev = torch.cat([radial_aev, angular_aev], dim=2)
+        return fullaev
+
+    def aev_to_output(self, aev, species):
+        """Compute output from aev
+
+        Parameters
+        ----------
+        aev : torch.Tensor
+            Pytorch tensor of shape (conformations, atoms, aev_length) storing
+            the computed AEVs.
+        species : list of string
+            List of string storing the species for each atom.
+
+        Returns
+        -------
+        torch.Tensor
+            Pytorch tensor of shape (conformations, output_length) for the
+            output of each conformation.
+        """
+        conformations = aev.shape[0]
+        atoms = len(species)
+        rev_species = species[::-1]
+        species_dedup = sorted(
+            set(species), key=self.aev_computer.species.index)
+        per_species_outputs = []
+        for s in species_dedup:
+            begin = species.index(s)
+            end = atoms - rev_species.index(s)
+            y = aev[:, begin:end, :].reshape(-1, self.aev_computer.aev_length)
+
+            def apply_model(suffix):
+                model_X = getattr(self, 'model_' + s + suffix)
+                return model_X(y)
+            ys = [apply_model(suffix) for suffix in self.suffixes]
+            y = sum(ys) / len(ys)
+            y = y.view(conformations, -1, self.output_length)
+            per_species_outputs.append(y)
+
+        per_species_outputs = torch.cat(per_species_outputs, dim=1)
+        molecule_output = self.reducer(per_species_outputs, dim=1)
+        return molecule_output
+
+    def compute_derivative(self, output, coordinates):
+        """Compute the gradient d(output)/d(coordinates)"""
+        # Since different conformations are independent, computing
+        # the derivatives of all outputs w.r.t. its own coordinate is
+        # equivalent to compute the derivative of the sum of all outputs
+        # w.r.t. all coordinates.
+        return torch.autograd.grad(output.sum(), coordinates,
+                                   create_graph=self.derivative_graph)[0]
+
+    def forward(self, coordinates, species):
+        """Feed forward
+
+        Parameters
+        ----------
+        coordinates : torch.Tensor
+            The pytorch tensor of shape (conformations, atoms, 3) storing
+            the coordinates of all atoms of all conformations.
+        species : list of string
+            List of string storing the species for each atom.
+
+        Returns
+        -------
+        torch.Tensor or (torch.Tensor, torch.Tensor)
+            If derivative is turned off, then this function will return a
+            pytorch tensor of shape (conformations, output_length) for the
+            output of each conformation.
+            If derivative is turned on, then this function will return a pair
+            of pytorch tensors where the first tensor is the output tensor as
+            when the derivative is off, and the second tensor is a tensor of
+            shape (conformation, atoms, 3) storing the d(output)/dR.
+        """
+        if not self.derivative:
+            coordinates = coordinates.detach()
+        else:
+            coordinates = torch.tensor(coordinates, requires_grad=True)
+        _coordinates, _species = self.aev_computer.sort_by_species(
+            coordinates, species)
+        aev = self.compute_aev(_coordinates, _species)
+        output = self.aev_to_output(aev, _species)
+        if not self.derivative:
+            return output
+        else:
+            derivative = self.compute_derivative(output, coordinates)
+            return output, derivative

torchani/models/custom.py

+from .ani_model import ANIModel
+
+
+class CustomModel(ANIModel):
+
+    def __init__(self, aev_computer, per_species, reducer,
+                 derivative=False, derivative_graph=False, benchmark=False):
+        """Custom single model, no ensemble
+
+        Parameters
+        ----------
+        per_species : dict
+            Dictionary with supported species as keys and objects of
+            `torch.nn.Model` as values, storing the model for each supported
+            species. These models will finally become `model_X` attributes.
+        reducer : function
+            The desired `reducer` attribute.
+        """
+        suffixes = ['']
+        output_length = None
+        models = {}
+        for i in per_species:
+            model_X = per_species[i]
+            if not hasattr(model_X, 'output_length'):
+                raise ValueError(
+                    '''atomic neural network must explicitly specify
+                    output length''')
+            elif output_length is None:
+                output_length = model_X.output_length
+            elif output_length != model_X.output_length:
+                raise ValueError(
+                    '''output length of each atomic neural network must
+                    match''')
+            setattr(self, 'model_' + i, model_X)
+        super(CustomModel, self).__init__(aev_computer, suffixes, reducer,
+                                          output_length, models, derivative,
+                                          derivative_graph, benchmark)

torchani/nn.py → torchani/models/neurochem_atomic_network.py

-from .aev_base import AEVComputer
-import torch
-import bz2
+from .. import _six  # noqa: F401
 import os
+import bz2
 import lark
-import struct
+import torch
 import math
-from .env import buildin_network_dir, buildin_model_prefix
-from .benchmarked import BenchmarkedModule
-
-# For python 2 compatibility
-if not hasattr(math, 'inf'):
-    math.inf = float('inf')
+import struct
 
 
-class PerSpeciesFromNeuroChem(torch.jit.ScriptModule):
-    """Subclass of `torch.nn.Module` for the per atom aev->y
-    transformation, loaded from NeuroChem network dir.
+class NeuroChemAtomicNetwork(torch.jit.ScriptModule):
+    """Per atom aev->y transformation, loaded from NeuroChem network dir.
 
     Attributes
     ----------
@@ -48,7 +41,7 @@ class PerSpeciesFromNeuroChem(torch.jit.ScriptModule):
             hyperparameters. The `.bparam` and `.wparam` must be
             in the same directory
         """
-        super(PerSpeciesFromNeuroChem, self).__init__()
+        super(NeuroChemAtomicNetwork, self).__init__()
 
         self.dtype = dtype
         self.device = device
@@ -284,292 +277,3 @@ class PerSpeciesFromNeuroChem(torch.jit.ScriptModule):
             output.
         """
         return self.get_activations(aev, math.inf)
-
-
-class ModelOnAEV(BenchmarkedModule):
-    """Subclass of `torch.nn.Module` for the [xyz]->[aev]->[per_atom_y]->y
-    pipeline.
-
-    Attributes
-    ----------
-    aev_computer : AEVComputer
-        The AEV computer.
-    output_length : int
-        The length of output vector
-    derivative : boolean
-        Whether to support computing the derivative w.r.t coordinates,
-        i.e. d(output)/dR
-    derivative_graph : boolean
-        Whether to generate a graph for the derivative. This would be required
-        only if the derivative is included as part of the loss function.
-    model_X : nn.Module
-        Model for species X. There should be one such attribute for each
-        supported species.
-    reducer : function
-        Function of (input, dim)->output that reduce the input tensor along the
-        given dimension to get an output tensor. This function will be called
-        with the per atom output tensor with internal shape as input, and
-        desired reduction dimension as dim, and should reduce the input into
-        the tensor containing desired output.
-    timers : dict
-        Dictionary storing the the benchmark result. It has the following keys:
-            aev : time spent on computing AEV.
-            nn : time spent on computing output from AEV.
-            derivative : time spend on computing derivative w.r.t. coordinates
-                after the outputs is given. This key is only available if
-                derivative computation is turned on.
-            forward : total time for the forward pass
-    """
-
-    def __init__(self, aev_computer, derivative=False, derivative_graph=False,
-                 benchmark=False, **kwargs):
-        """Initialize object from manual setup or from NeuroChem network
-        directory.
-
-        The caller must set either `from_nc` in order to load from NeuroChem
-        network directory, or set `per_species` and `reducer`.
-
-        Parameters
-        ----------
-        aev_computer : AEVComputer
-            The AEV computer.
-        derivative : boolean
-            Whether to support computing the derivative w.r.t coordinates,
-            i.e. d(output)/dR
-        derivative_graph : boolean
-            Whether to generate a graph for the derivative. This would be
-            required only if the derivative is included as part of the loss
-            function. This argument must be set to False if `derivative` is
-            set to False.
-        benchmark : boolean
-            Whether to enable benchmarking
-
-        Other Parameters
-        ----------------
-        from_nc : string
-            Path to the NeuroChem network directory. If this parameter is set,
-            then `per_species` and `reducer` should not be set. If set to
-            `None`, then the network ship with torchani will be used.
-        ensemble : int
-            Number of models in the model ensemble. If this is not set, then
-            `from_nc` would refer to the directory storing the model. If set to
-            a number, then `from_nc` would refer to the prefix of directories.
-        per_species : dict
-            Dictionary with supported species as keys and objects of
-            `torch.nn.Model` as values, storing the model for each supported
-            species. These models will finally become `model_X` attributes.
-        reducer : function
-            The desired `reducer` attribute.
-
-        Raises
-        ------
-        ValueError
-            If `from_nc`, `per_species`, and `reducer` are not properly set.
-        """
-
-        super(ModelOnAEV, self).__init__(benchmark)
-        self.derivative = derivative
-        self.output_length = None
-        if not derivative and derivative_graph:
-            raise ValueError(
-                '''ModelOnAEV: can not create graph for derivative if the
-                computation of derivative is turned off''')
-        self.derivative_graph = derivative_graph
-
-        if benchmark:
-            self.compute_aev = self._enable_benchmark(self.compute_aev, 'aev')
-            self.aev_to_output = self._enable_benchmark(
-                self.aev_to_output, 'nn')
-            if derivative:
-                self.compute_derivative = self._enable_benchmark(
-                    self.compute_derivative, 'derivative')
-            self.forward = self._enable_benchmark(self.forward, 'forward')
-
-        if not isinstance(aev_computer, AEVComputer):
-            raise TypeError(
-                "ModelOnAEV: aev_computer must be a subclass of AEVComputer")
-        self.aev_computer = aev_computer
-
-        if 'from_nc' in kwargs and 'per_species' not in kwargs and \
-           'reducer' not in kwargs:
-            if 'ensemble' not in kwargs:
-                if kwargs['from_nc'] is None:
-                    kwargs['from_nc'] = buildin_network_dir
-                network_dirs = [kwargs['from_nc']]
-                self.suffixes = ['']
-            else:
-                if kwargs['from_nc'] is None:
-                    kwargs['from_nc'] = buildin_model_prefix
-                network_prefix = kwargs['from_nc']
-                network_dirs = []
-                self.suffixes = []
-                for i in range(kwargs['ensemble']):
-                    suffix = '{}'.format(i)
-                    network_dir = os.path.join(
-                        network_prefix+suffix, 'networks')
-                    network_dirs.append(network_dir)
-                    self.suffixes.append(suffix)
-
-            self.reducer = torch.sum
-            for network_dir, suffix in zip(network_dirs, self.suffixes):
-                for i in self.aev_computer.species:
-                    filename = os.path.join(
-                        network_dir, 'ANN-{}.nnf'.format(i))
-                    model_X = PerSpeciesFromNeuroChem(
-                        self.aev_computer.dtype, self.aev_computer.device,
-                        filename)
-                    if self.output_length is None:
-                        self.output_length = model_X.output_length
-                    elif self.output_length != model_X.output_length:
-                        raise ValueError(
-                            '''output length of each atomic neural networt
-                            must match''')
-                    setattr(self, 'model_' + i + suffix, model_X)
-        elif 'from_nc' not in kwargs and 'per_species' in kwargs and \
-             'reducer' in kwargs:
-            self.suffixes = ['']
-            per_species = kwargs['per_species']
-            for i in per_species:
-                model_X = per_species[i]
-                if not hasattr(model_X, 'output_length'):
-                    raise ValueError(
-                        '''atomic neural network must explicitly specify
-                        output length''')
-                elif self.output_length is None:
-                    self.output_length = model_X.output_length
-                elif self.output_length != model_X.output_length:
-                    raise ValueError(
-                        '''output length of each atomic neural network must
-                        match''')
-                setattr(self, 'model_' + i, model_X)
-            self.reducer = kwargs['reducer']
-        else:
-            raise ValueError(
-                'ModelOnAEV: bad arguments when initializing ModelOnAEV')
-
-        if derivative and self.output_length != 1:
-            raise ValueError(
-                'derivative can only be computed for output length 1')
-
-    def compute_aev(self, coordinates, species):
-        """Compute full AEV
-
-        Parameters
-        ----------
-        coordinates : torch.Tensor
-            The pytorch tensor of shape (conformations, atoms, 3) storing
-            the coordinates of all atoms of all conformations.
-        species : list of string
-            List of string storing the species for each atom.
-
-        Returns
-        -------
-        torch.Tensor
-            Pytorch tensor of shape (conformations, atoms, aev_length) storing
-            the computed AEVs.
-        """
-        radial_aev, angular_aev = self.aev_computer(coordinates, species)
-        fullaev = torch.cat([radial_aev, angular_aev], dim=2)
-        return fullaev
-
-    def aev_to_output(self, aev, species):
-        """Compute output from aev
-
-        Parameters
-        ----------
-        aev : torch.Tensor
-            Pytorch tensor of shape (conformations, atoms, aev_length) storing
-            the computed AEVs.
-        species : list of string
-            List of string storing the species for each atom.
-
-        Returns
-        -------
-        torch.Tensor
-            Pytorch tensor of shape (conformations, output_length) for the
-            output of each conformation.
-        """
-        conformations = aev.shape[0]
-        atoms = len(species)
-        rev_species = species[::-1]
-        species_dedup = sorted(
-            set(species), key=self.aev_computer.species.index)
-        per_species_outputs = []
-        for s in species_dedup:
-            begin = species.index(s)
-            end = atoms - rev_species.index(s)
-            y = aev[:, begin:end, :].contiguous(
-            ).view(-1, self.aev_computer.aev_length)
-
-            def apply_model(suffix):
-                model_X = getattr(self, 'model_' + s + suffix)
-                return model_X(y)
-            ys = [apply_model(suffix) for suffix in self.suffixes]
-            y = sum(ys) / len(ys)
-            y = y.view(conformations, -1, self.output_length)
-            per_species_outputs.append(y)
-
-        per_species_outputs = torch.cat(per_species_outputs, dim=1)
-        molecule_output = self.reducer(per_species_outputs, dim=1)
-        return molecule_output
-
-    def compute_derivative(self, output, coordinates):
-        """Compute the gradient d(output)/d(coordinates)"""
-        # Since different conformations are independent, computing
-        # the derivatives of all outputs w.r.t. its own coordinate is
-        # equivalent to compute the derivative of the sum of all outputs
-        # w.r.t. all coordinates.
-        return torch.autograd.grad(output.sum(), coordinates,
-                                   create_graph=self.derivative_graph)[0]
-
-    def forward(self, coordinates, species):
-        """Feed forward
-
-        Parameters
-        ----------
-        coordinates : torch.Tensor
-            The pytorch tensor of shape (conformations, atoms, 3) storing
-            the coordinates of all atoms of all conformations.
-        species : list of string
-            List of string storing the species for each atom.
-
-        Returns
-        -------
-        torch.Tensor or (torch.Tensor, torch.Tensor)
-            If derivative is turned off, then this function will return a
-            pytorch tensor of shape (conformations, output_length) for the
-            output of each conformation.
-            If derivative is turned on, then this function will return a pair
-            of pytorch tensors where the first tensor is the output tensor as
-            when the derivative is off, and the second tensor is a tensor of
-            shape (conformation, atoms, 3) storing the d(output)/dR.
-        """
-        if not self.derivative:
-            coordinates = coordinates.detach()
-        else:
-            coordinates = torch.tensor(coordinates, requires_grad=True)
-        _coordinates, _species = self.aev_computer.sort_by_species(
-            coordinates, species)
-        aev = self.compute_aev(_coordinates, _species)
-        output = self.aev_to_output(aev, _species)
-        if not self.derivative:
-            return output
-        else:
-            derivative = self.compute_derivative(output, coordinates)
-            return output, derivative
-
-    def export_onnx(self, dirname):
-        """Export atomic networks into onnx format
-
-        Parameters
-        ----------
-        dirname : string
-            Name of the directory to store exported networks.
-        """
-
-        aev_length = self.aev_computer.aev_length
-        dummy_aev = torch.zeros(1, aev_length)
-        for s in self.aev_computer.species:
-            nn_onnx = os.path.join(dirname, '{}.proto'.format(s))
-            model_X = getattr(self, 'model_' + s)
-            torch.onnx.export(model_X, dummy_aev, nn_onnx)

torchani/models/neurochem_nnp.py

+import os
+import torch
+from .ani_model import ANIModel
+from .neurochem_atomic_network import NeuroChemAtomicNetwork
+from ..env import buildin_network_dir, buildin_model_prefix, buildin_ensemble
+
+
+class NeuroChemNNP(ANIModel):
+
+    def __init__(self, aev_computer, from_=None, ensemble=False,
+                 derivative=False, derivative_graph=False, benchmark=False):
+        """If from_=None then ensemble must be a boolean. If ensemble=False,
+            then use buildin network0, else use buildin network ensemble.
+        If from_ != None, ensemble must be either False or an integer
+            specifying the number of networks in the ensemble.
+        """
+        if from_ is None:
+            if not isinstance(ensemble, bool):
+                raise TypeError('ensemble must be boolean')
+            if ensemble:
+                from_ = buildin_model_prefix
+                ensemble = buildin_ensemble
+            else:
+                from_ = buildin_network_dir
+        else:
+            if not (ensemble is False or isinstance(ensemble, int)):
+                raise ValueError('invalid argument ensemble')
+
+        if ensemble is False:
+            network_dirs = [from_]
+            suffixes = ['']
+        else:
+            assert isinstance(ensemble, int)
+            network_prefix = from_
+            network_dirs = []
+            suffixes = []
+            for i in range(ensemble):
+                suffix = '{}'.format(i)
+                network_dir = os.path.join(
+                    network_prefix+suffix, 'networks')
+                network_dirs.append(network_dir)
+                suffixes.append(suffix)
+
+        reducer = torch.sum
+
+        models = {}
+        output_length = None
+        for network_dir, suffix in zip(network_dirs, suffixes):
+            for i in aev_computer.species:
+                filename = os.path.join(
+                    network_dir, 'ANN-{}.nnf'.format(i))
+                model_X = NeuroChemAtomicNetwork(
+                    aev_computer.dtype, aev_computer.device,
+                    filename)
+                if output_length is None:
+                    output_length = model_X.output_length
+                elif output_length != model_X.output_length:
+                    raise ValueError(
+                        '''output length of each atomic neural networt
+                        must match''')
+                models['model_' + i + suffix] = model_X
+        super(NeuroChemNNP, self).__init__(aev_computer, suffixes, reducer,
+                                           output_length, models, derivative,
+                                           derivative_graph, benchmark)