Merge pull request #56 from ziatdinovmax/master

Add ensemble DKL and notebook with example

atomai/models/dklgp/dklgpr.py

@@ -82,7 +82,8 @@ def fit(self, X: Union[torch.Tensor, np.ndarray],
 
         Keyword Args:
             feature_extractor:
-                (Optional) Custom neural network for feature extractor
+                (Optional) Custom neural network for feature extractor.
+                Must take input/feature dims and embedding dims as its arguments.
             freeze_weights:
                 Freezes weights of feature extractor, that is, they are not
                 passed to the optimizer. Used for a transfer learning.
@@ -91,6 +92,45 @@ def fit(self, X: Union[torch.Tensor, np.ndarray],
         """
         _ = self.run(X, y, training_cycles, **kwargs)
 
+    def fit_ensemble(self, X: Union[torch.Tensor, np.ndarray],
+                     y: Union[torch.Tensor, np.ndarray],
+                     training_cycles: int = 1,
+                     n_models: int = 5,
+                     **kwargs: Union[Type[torch.nn.Module], bool, float]
+                     ) -> None:
+        """
+        Initializes and trains an ensemble of deep kernel GP model
+
+        Args:
+            X: Input training data (aka features) of N x input_dim dimensions
+            y: Output targets of batch_size x N or N (if batch_size=1) dimensions
+            training_cycles: Number of training epochs
+            n_models: Number of models in ensemble
+
+        Keyword Args:
+            feature_extractor:
+                (Optional) Custom neural network for feature extractor.
+                Must take input/feature dims and embedding dims as its arguments.
+            freeze_weights:
+                Freezes weights of feature extractor, that is, they are not
+                passed to the optimizer. Used for a transfer learning.
+            lr: learning rate (Default: 0.01)
+            print_loss: print loss at every n-th training cycle (epoch)
+        """
+        if y.ndim == 1:
+            y = y[None]
+        if y.shape[0] > 1:
+            raise NotImplementedError(
+                "The ensemble training is currently supported only for scalar targets")
+        y = y.repeat(n_models, 0) if isinstance(y, np.ndarray) else y.repeat(n_models, 1)
+        if self.correlated_output:
+            msg = ("Replacing shared independent embedding space with" +
+                   " {} independent ones").format(n_models)
+            warnings.warn(msg)
+            self.correlated_output = False
+        self.ensemble = True
+        _ = self.run(X, y, training_cycles, **kwargs)
+
     def _compute_posterior(self, X: torch.Tensor) -> Union[mvn_, List[mvn_]]:
         """
         Computes the posterior over model outputs at the provided points (X).
@@ -194,7 +234,7 @@ def embed(self, x_new: Union[torch.Tensor, np.ndarray],
         x_new, _ = self.set_data(x_new, device='cpu')
         data_loader = init_dataloader(x_new, shuffle=False, **kwargs)
         embeded = torch.cat([self._embed(x.to(self.device)) for (x,) in data_loader], 0)
-        if not self.correlated_output:
+        if not self.correlated_output and not self.ensemble:
             embeded = embeded.permute(-1, 0, 1)
         return embeded.numpy()
 

atomai/trainers/gptrainer.py

@@ -52,6 +52,7 @@ def __init__(self,
         self.correlated_output = shared_embedding_space
         self.gp_model = None
         self.likelihood = None
+        self.ensemble = False
         self.compiled = False
         self.train_loss = []
 
@@ -92,6 +93,8 @@ def compile_multi_model_trainer(self,
         the number of neural networks is equal to the number of Gaussian
         processes. For example, if the outputs are spectra of length 128,
         one will have 128 neural networks and 128 GPs trained in parallel.
+        It can be also used for training an ensembles of models for the same
+        scalar output.
         """
         if self.correlated_output:
             raise NotImplementedError(
@@ -101,16 +104,21 @@ def compile_multi_model_trainer(self,
         if y.shape[0] < 2:
             raise ValueError("The training targets must be vector-valued (d >1)")
         input_dim, embedim = self.dimdict["input_dim"], self.dimdict["embedim"]
-        feature_extractor = kwargs.get("feature_extractor")
-        if feature_extractor is None:
-            feature_extractor = fcFeatureExtractor(input_dim, embedim)
+        feature_net = kwargs.get("feature_extractor", fcFeatureExtractor)
         freeze_weights = kwargs.get("freeze_weights", False)
-        if freeze_weights:
-            for p in feature_extractor.parameters():
-                p.requires_grad = False
+        if not self.ensemble:
+            feature_extractor = feature_net(input_dim, embedim)
+            if freeze_weights:
+                for p in feature_extractor.parameters():
+                    p.requires_grad = False
         list_of_models = []
         list_of_likelihoods = []
         for i in range(y.shape[0]):
+            if self.ensemble:  # different initilization for each model
+                feature_extractor = feature_net(input_dim, embedim)
+                if freeze_weights:
+                    for p in feature_extractor.parameters():
+                        p.requires_grad = False
             model_i = GPRegressionModel(
                 X, y[i:i+1],
                 gpytorch.likelihoods.GaussianLikelihood(batch_shape=torch.Size([1])),
@@ -154,7 +162,8 @@ def compile_trainer(self, X: Union[torch.Tensor, np.ndarray],
 
         Keyword Args:
             feature_extractor:
-                (Optional) Custom neural network for feature extractor
+                (Optional) Custom neural network for feature extractor.
+                Must take input/feature dims and embedding dims as its arguments.
             grid_size:
                 Grid size for structured kernel interpolation (Default: 50)
             freeze_weights:
@@ -168,9 +177,8 @@ def compile_trainer(self, X: Union[torch.Tensor, np.ndarray],
                 "use compile_multi_model_trainer(*args, **kwargs)")
         X, y = self.set_data(X, y)
         input_dim, embedim = self.dimdict["input_dim"], self.dimdict["embedim"]
-        feature_extractor = kwargs.get("feature_extractor")
-        if feature_extractor is None:
-            feature_extractor = fcFeatureExtractor(input_dim, embedim)
+        feature_net = kwargs.get("feature_extractor", fcFeatureExtractor)
+        feature_extractor = feature_net(input_dim, embedim)
         freeze_weights = kwargs.get("freeze_weights", False)
         if freeze_weights:
             for p in feature_extractor.parameters():