Robust classifier (deepchem#4159)

* Robust classifier

* Adds more tasks to the forward test for architecture similarity

* Resolves conflict and passes lint tests

* Fixes input label processing using default generator

* Updates weights files with more appropriate ones with more tasks

* Formats for yapf and flake8

* Added tests for classification

* Fixes reload test

* Added PyTorch robustmultitask classifier model to rst

* Added docstring to default generator and logs training per epoch

* Removes logger and docstring

* Adds docstring to default generator and logs epoch training

deepchem/models/torch_models/__init__.py

@@ -34,7 +34,7 @@
 from deepchem.models.torch_models.unet import UNet, UNetModel
 from deepchem.models.torch_models.graphconvmodel import _GraphConvTorchModel, GraphConvModel
 from deepchem.models.torch_models.smiles2vec import Smiles2Vec, Smiles2VecModel
-from deepchem.models.torch_models.robust_multitask import RobustMultitask
+from deepchem.models.torch_models.robust_multitask import RobustMultitask, RobustMultitaskClassifier
 from deepchem.models.torch_models.inceptionv3 import InceptionV3Model, InceptionA, InceptionB, InceptionC, InceptionD, InceptionE, InceptionAux, BasicConv2d
 try:
     from deepchem.models.torch_models.dmpnn import DMPNN, DMPNNModel

deepchem/models/torch_models/robust_multitask.py

@@ -4,7 +4,11 @@
 import logging
 from typing import List, Tuple, Callable, Literal, Union
 from typing import Sequence as SequenceCollection
+from deepchem.models.torch_models.torch_model import TorchModel
+from deepchem.models import losses
 from deepchem.utils.typing import OneOrMany, ActivationFn
+from deepchem.metrics import to_one_hot
+import datetime
 
 logger = logging.getLogger(__name__)
 
@@ -205,7 +209,6 @@ def forward(
             task_outputs.append(task_output)
 
         output = torch.stack(task_outputs, dim=1)
-
         if self.mode == 'classification':
             if self.n_tasks == 1:
                 logits = output.view(-1, self.n_classes)
@@ -251,3 +254,159 @@ def _get_activation_class(self, activation_name: ActivationFn) -> Callable:
             raise ValueError(
                 f"Invalid activation function: {activation_name}. Only activations of type nn.Module"
             )
+
+
+class RobustMultitaskClassifier(TorchModel):
+    """
+    Implements a neural network for robust multitasking.
+    The key idea of this model is to have bypass layers that feed
+    directly from features to task output. This might provide some
+    flexibility toroute around challenges in multitasking with
+    destructive interference.
+    References
+    ----------
+    This technique was introduced in [1]_
+    .. [1] Ramsundar, Bharath, et al. "Is multitask deep learning practical for pharma?." Journal of chemical information and modeling 57.8 (2017): 2068-2076.
+    """
+
+    def __init__(self,
+                 n_tasks: int,
+                 n_features: int,
+                 layer_sizes: SequenceCollection[int] = [1000],
+                 weight_init_stddevs: OneOrMany[float] = 0.02,
+                 bias_init_consts: OneOrMany[float] = 1.0,
+                 weight_decay_penalty: float = 0.0,
+                 weight_decay_penalty_type: Literal['l1', 'l2'] = "l2",
+                 dropouts: OneOrMany[float] = 0.5,
+                 activation_fns: OneOrMany[ActivationFn] = nn.ReLU(),
+                 n_classes: int = 2,
+                 bypass_layer_sizes: SequenceCollection[int] = [100],
+                 bypass_weight_init_stddevs: OneOrMany[float] = [0.02],
+                 bypass_bias_init_consts: OneOrMany[float] = [1.0],
+                 bypass_dropouts: OneOrMany[float] = [0.5],
+                 **kwargs):
+        """
+        Parameters
+        ----------
+        n_tasks: int
+            number of tasks
+        n_features: int
+            number of features
+        layer_sizes: list
+            the size of each dense layer in the network.  The length of this list determines the number of layers.
+        weight_init_stddevs: list or float
+            the standard deviation of the distribution to use for weight initialization of each layer.  The length
+            of this list should equal len(layer_sizes).  Alternatively this may be a single value instead of a list,
+            in which case the same value is used for every layer.
+        bias_init_consts: list or loat
+            the value to initialize the biases in each layer to.  The length of this list should equal len(layer_sizes).
+            Alternatively this may be a single value instead of a list, in which case the same value is used for every layer.
+        weight_decay_penalty: float
+            the magnitude of the weight decay penalty to use
+        weight_decay_penalty_type: str
+            the type of penalty to use for weight decay, either 'l1' or 'l2'
+        dropouts: list or float
+            the dropout probablity to use for each layer.  The length of this list should equal len(layer_sizes).
+            Alternatively this may be a single value instead of a list, in which case the same value is used for every layer.
+        activation_fns: list or object
+            the Tensorflow activation function to apply to each layer.  The length of this list should equal
+            len(layer_sizes).  Alternatively this may be a single value instead of a list, in which case the
+            same value is used for every layer.
+        n_classes: int
+            the number of classes
+        bypass_layer_sizes: list
+            the size of each dense layer in the bypass network. The length of this list determines the number of bypass layers.
+        bypass_weight_init_stddevs: list or float
+            the standard deviation of the distribution to use for weight initialization of bypass layers.
+            same requirements as weight_init_stddevs
+        bypass_bias_init_consts: list or float
+            the value to initialize the biases in bypass layers
+            same requirements as bias_init_consts
+        bypass_dropouts: list or float
+            the dropout probablity to use for bypass layers.
+            same requirements as dropouts
+        """
+        if not isinstance(activation_fns, nn.Module):
+            logger.warning(
+                "Warning: Activation functions should be of type nn.Module. Using default activation function: ReLU."
+            )
+            activation_fns = nn.ReLU()
+
+        # The labels are one-hot encoded.
+        loss = losses.SoftmaxCrossEntropy()
+        output_types = ['prediction', 'loss']
+        self.n_classes = n_classes
+        self.n_tasks = n_tasks
+
+        model = RobustMultitask(
+            n_tasks=self.n_tasks,
+            n_features=n_features,
+            layer_sizes=layer_sizes,
+            mode='classification',
+            weight_init_stddevs=weight_init_stddevs,
+            bias_init_consts=bias_init_consts,
+            weight_decay_penalty=weight_decay_penalty,
+            weight_decay_penalty_type=weight_decay_penalty_type,
+            activation_fns=activation_fns,
+            dropouts=dropouts,
+            n_classes=self.n_classes,
+            bypass_layer_sizes=bypass_layer_sizes,
+            bypass_weight_init_stddevs=bypass_weight_init_stddevs,
+            bypass_bias_init_consts=bypass_bias_init_consts,
+            bypass_dropouts=bypass_dropouts)
+        self.activation_fns = model.activation_fns
+        self.dropouts = model.dropouts
+        self.shared_layers = model.shared_layers
+        self.bypass_layers = model.bypass_layers
+        self.output_layers = model.output_layers
+
+        super(RobustMultitaskClassifier,
+              self).__init__(model,
+                             loss,
+                             output_types=output_types,
+                             regularization_loss=model.regularization_loss,
+                             **kwargs)
+
+    def default_generator(self,
+                          dataset,
+                          epochs=1,
+                          mode='fit',
+                          deterministic=True,
+                          pad_batches=True):
+        """Create a generator that iterates batches for a dataset.
+
+        Subclasses may override this method to customize how model inputs are
+        generated from the data.
+
+        Parameters
+        ----------
+        dataset: Dataset
+            the data to iterate
+        epochs: int
+            the number of times to iterate over the full dataset
+        mode: str
+            allowed values are 'fit' (called during training), 'predict' (called
+            during prediction), and 'uncertainty' (called during uncertainty
+            prediction)
+        deterministic: bool
+            whether to iterate over the dataset in order, or randomly shuffle the
+            data for each epoch
+        pad_batches: bool
+            whether to pad each batch up to this model's preferred batch size
+
+        Returns
+        -------
+        a generator that iterates batches, each represented as a tuple of lists:
+        ([inputs], [outputs], [weights])
+        """
+        for epoch in range(epochs):
+            logger.info("Starting training for epoch %d at %s" %
+                        (epoch, datetime.datetime.now().ctime()))
+            for (X_b, y_b, w_b,
+                 ids_b) in dataset.iterbatches(batch_size=self.batch_size,
+                                               deterministic=deterministic,
+                                               pad_batches=pad_batches):
+                if y_b is not None:
+                    y_b = to_one_hot(y_b.flatten(), self.n_classes).reshape(
+                        -1, self.n_tasks, self.n_classes)
+                yield ([X_b], [y_b], [w_b])