Implemented TorchModel.load_from_pretrained()

import os

import unittest

import deepchem as dc

import numpy as np

from deepchem.models.losses import L2Loss

from deepchem.feat.mol_graphs import ConvMol

try:

  import torch

  has_pytorch = True

except:

  has_pytorch = False

class MLP(dc.models.TorchModel):

  def __init__(self, n_tasks=1, feature_dim=100, hidden_layer_size=64,

               **kwargs):

    pytorch_model = torch.nn.Sequential(

        torch.nn.Linear(feature_dim, hidden_layer_size), torch.nn.ReLU(),

        torch.nn.Linear(hidden_layer_size, n_tasks), torch.nn.Sigmoid())

    loss = dc.models.losses.BinaryCrossEntropy()

    super(MLP, self).__init__(model=pytorch_model, loss=loss, **kwargs)

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

class TestPretrainedTorch(unittest.TestCase):

  def setUp(self):

    self.feature_dim = 2

    self.hidden_layer_size = 10

    data_points = 10

    X = np.random.randn(data_points, self.feature_dim)

    y = (X[:, 0] > X[:, 1]).astype(np.float32)

    self.dataset = dc.data.NumpyDataset(X, y)

  def test_load_from_pretrained(self):

    """Tests loading pretrained model."""

    source_model = MLP(

        hidden_layer_size=self.hidden_layer_size,

        feature_dim=self.feature_dim,

        batch_size=10)

    source_model.fit(self.dataset, nb_epoch=1000, checkpoint_interval=0)

    dest_model = MLP(

        feature_dim=self.feature_dim,

        hidden_layer_size=self.hidden_layer_size,

        n_tasks=10)

    assignment_map = dict()

    value_map = dict()

    source_vars = list(source_model.model.parameters())

    dest_vars = list(dest_model.model.parameters())[:-2]

    for idx, dest_var in enumerate(dest_vars):

      source_var = source_vars[idx]

      assignment_map[source_var] = dest_var

      value_map[source_var] = source_var.detach().numpy()

    dest_model.load_from_pretrained(

        source_model=source_model,

        assignment_map=assignment_map,

        value_map=value_map)

    for source_var, dest_var in assignment_map.items():

      source_val = source_var.detach().numpy()

      dest_val = dest_var.detach().numpy()

      np.testing.assert_array_almost_equal(source_val, dest_val)

  def test_restore_equivalency(self):

    """Test for restore based pretrained model loading."""

    source_model = MLP(

        feature_dim=self.feature_dim, hidden_layer_size=self.hidden_layer_size)

    source_model.fit(self.dataset, nb_epoch=1000)

    dest_model = MLP(

        feature_dim=self.feature_dim, hidden_layer_size=self.hidden_layer_size)

    dest_model.load_from_pretrained(

        source_model=source_model,

        assignment_map=None,

        value_map=None,

        model_dir=None,

        include_top=True)

    predictions = np.squeeze(dest_model.predict_on_batch(self.dataset.X))

    np.testing.assert_array_almost_equal(self.dataset.y, np.round(predictions))

try:

  import torch

  import torch.nn.functional as F

  has_pytorch = True

except:

  has_pytorch = False

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_overfit_subclass_model():

  """Test fitting a TorchModel defined by subclassing Module."""

  import torch.nn.functional as F

  n_data_points = 10

  n_features = 2

  np.random.seed(1234)

  assert scores[metric.name] > 0.9

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_overfit_sequential_model():

  """Test fitting a TorchModel defined as a sequential model."""

  n_data_points = 10

  assert scores[metric.name] > 0.9

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_fit_use_all_losses():

  """Test fitting a TorchModel and getting a loss curve back."""

  n_data_points = 10

  assert np.count_nonzero(np.array(losses)) == 100

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_fit_on_batch():

  """Test fitting a TorchModel to individual batches."""

  n_data_points = 10

  assert scores[metric.name] > 0.9

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_checkpointing():

  """Test loading and saving checkpoints with TorchModel."""

  # Create two models using the same model directory.

  assert np.array_equal(y1, y4)

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_fit_restore():

  """Test specifying restore=True when calling fit()."""

  n_data_points = 10

  assert np.array_equal(y, np.round(prediction))

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_uncertainty():

  """Test estimating uncertainty a TorchModel."""

  n_samples = 30

      self.log_var = torch.nn.Linear(200, n_features)

    def forward(self, inputs):

      import torch.nn.functional as F

      x, use_dropout = inputs

      x = self.hidden(x)

      if use_dropout:

  assert noise < np.mean(std) < 1.0

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_saliency_mapping():

  """Test computing a saliency map."""

  n_tasks = 3

    assert np.allclose(pred1[task], (pred2 + norm * delta)[task])

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_saliency_shapes():

  """Test computing saliency maps for multiple outputs with multiple dimensions."""

  assert file_size > 0

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_fit_variables():

  """Test training a subset of the variables in a model."""

  assert np.allclose(vars[1], 0.5)

@unittest.skipIf(not has_pytorch, 'PyTorch is not installed')

def test_fit_loss():

  """Test specifying a different loss function when calling fit()."""

    restore: bool

      if True, restore the model from the most recent checkpoint and continue training

      from there.  If False, retrain the model from scratch.

    variables: list of tf.Variable

    variables: list of torch.nn.Parameter

      the variables to train.  If None (the default), all trainable variables in

      the model are used.

    loss: function

    restore: bool

      if True, restore the model from the most recent checkpoint and continue training

      from there.  If False, retrain the model from scratch.

    variables: list of tf.Variable

    variables: list of torch.nn.Parameter

      the variables to train.  If None (the default), all trainable variables in

      the model are used.

    loss: function

      the labels for the batch

    w: ndarray

      the weights for the batch

    variables: list of tf.Variable

    variables: list of torch.nn.Parameter

      the variables to train.  If None (the default), all trainable variables in

      the model are used.

    loss: function

    self._ensure_built()

    X, _, _ = self._prepare_batch(([X], None, None))

    # Use a GradientTape to compute gradients.

    # Compute the gradients.

    X = torch.Tensor(X[0])

    X.requires_grad_(True)

      outputs = [outputs]

    final_result = []

    for output in outputs:

      print(output.shape)

      output_shape = tuple(output.shape[1:])

      output = output.reshape([-1])

      result = []

    """

    self._ensure_built()

    if checkpoint is None:

      checkpoints = self.get_checkpoints(model_dir)

      checkpoints = sorted(self.get_checkpoints(model_dir))

      if len(checkpoints) == 0:

        raise ValueError('No checkpoint found')

      checkpoint = checkpoints[-1]

      checkpoint = checkpoints[0]

    data = torch.load(checkpoint)

    self.model.load_state_dict(data['model_state_dict'])

    self._pytorch_optimizer.load_state_dict(data['optimizer_state_dict'])

    """Get the number of steps of fitting that have been performed."""

    return self._global_step

  # def _create_assignment_map(self,

  #                            source_model: "TorchModel",

  #                            include_top: bool = True,

  #                            **kwargs) -> Dict[Any, Any]:

  #   """

  #   Creates a default assignment map between variables of source and current model.

  #   This is used only when a custom assignment map is missing. This assumes the

  #   model is made of different layers followed by a dense layer for mapping to

  #   output tasks. include_top is used to control whether or not the final dense

  #   layer is used. The default assignment map is useful in cases where the type

  #   of task is different (classification vs regression) and/or number of tasks.

  #

  #   Parameters

  #   ----------

  #   source_model: dc.models.TorchModel

  #       Source model to copy variable values from.

  #   include_top: bool, default True

  #       if true, copies the last dense layer

  #   """

  #   assignment_map: Dict[Any, Any] = {}

  #   source_vars = source_model.model.trainable_variables

  #   dest_vars = self.model.trainable_variables

  #

  #   if not include_top:

  #     source_vars = source_vars[:-2]

  #     dest_vars = dest_vars[:-2]

  #

  #   for source_var, dest_var in zip(source_vars, dest_vars):

  #     assignment_map[source_var.ref()] = dest_var

  #

  #   return assignment_map

  #

  # def _create_value_map(self, source_model: "TorchModel",

  #                       **kwargs) -> Dict[Any, Any]:

  #   """

  #   Creates a value map between variables in the source model and their

  #   current values. This is used only when a custom value map is missing, and

  #   assumes the restore method has been called under self.session.

  #

  #   Parameters

  #   ----------

  #   source_model: dc.models.TorchModel

  #       Source model to create value map from

  #   """

  #   value_map: Dict[Any, Any] = {}

  #   source_vars = source_model.model.trainable_variables

  #

  #   for source_var in source_vars:

  #     value_map[source_var.ref()] = source_var.numpy()

  #

  #   return value_map

  #

  # def load_from_pretrained(self,

  #                          source_model: "TorchModel",

  #                          assignment_map: Optional[Dict[Any, Any]] = None,

  #                          value_map: Optional[Dict[Any, Any]] = None,

  #                          checkpoint: Optional[str] = None,

  #                          model_dir: Optional[str] = None,

  #                          include_top: bool = True,

  #                          inputs: Optional[Sequence[Any]] = None,

  #                          **kwargs) -> None:

  #   """Copies variable values from a pretrained model. `source_model` can either

  #   be a pretrained model or a model with the same architecture. `value_map`

  #   is a variable-value dictionary. If no `value_map` is provided, the variable

  #   values are restored to the `source_model` from a checkpoint and a default

  #   `value_map` is created. `assignment_map` is a dictionary mapping variables

  #   from the `source_model` to the current model. If no `assignment_map` is

  #   provided, one is made from scratch and assumes the model is composed of

  #   several different layers, with the final one being a dense layer. include_top

  #   is used to control whether or not the final dense layer is used. The default

  #   assignment map is useful in cases where the type of task is different

  #   (classification vs regression) and/or number of tasks in the setting.

  #

  #   Parameters

  #   ----------

  #   source_model: dc.TorchModel, required

  #     source_model can either be the pretrained model or a dc.TorchModel with

  #     the same architecture as the pretrained model. It is used to restore from

  #     a checkpoint, if value_map is None and to create a default assignment map

  #     if assignment_map is None

  #   assignment_map: Dict, default None

  #     Dictionary mapping the source_model variables and current model variables

  #   value_map: Dict, default None

  #     Dictionary containing source_model trainable variables mapped to numpy

  #     arrays. If value_map is None, the values are restored and a default

  #     variable map is created using the restored values

  #   checkpoint: str, default None

  #     the path to the checkpoint file to load.  If this is None, the most recent

  #     checkpoint will be chosen automatically.  Call get_checkpoints() to get a

  #     list of all available checkpoints

  #   model_dir: str, default None

  #     Restore model from custom model directory if needed

  #   include_top: bool, default True

  #       if True, copies the weights and bias associated with the final dense

  #       layer. Used only when assignment map is None

  #   inputs: List, input tensors for model

  #       if not None, then the weights are built for both the source and self.

  #       This option is useful only for models that are built by

  #       subclassing torch.nn.Module, and not using the functional API by tf.keras

  #   """

  #   if inputs is not None:

  #     # Ensure weights for both models are built.

  #     source_model.model(inputs)

  #     self.model(inputs)

  #

  #   self._ensure_built()

  #   if value_map is None:

  #     logger.info(

  #         "No value map provided. Creating default value map from restored model."

  #     )

  #     source_model.restore(model_dir=model_dir, checkpoint=checkpoint)

  #     value_map = self._create_value_map(source_model=source_model)

  #

  #   if assignment_map is None:

  #     logger.info("No assignment map provided. Creating custom assignment map.")

  #     assignment_map = self._create_assignment_map(

  #         source_model=source_model, include_top=include_top)

  #

  #   for source_var, dest_var in assignment_map.items():

  #     assert source_var.deref().shape == dest_var.shape

  #     dest_var.assign(value_map[source_var])

  def _create_assignment_map(self,

                             source_model: "TorchModel",

                             include_top: bool = True,

                             **kwargs) -> Dict[Any, Any]:

    """

    Creates a default assignment map between parameters of source and current model.

    This is used only when a custom assignment map is missing. This assumes the

    model is made of different layers followed by a dense layer for mapping to

    output tasks. include_top is used to control whether or not the final dense

    layer is used. The default assignment map is useful in cases where the type

    of task is different (classification vs regression) and/or number of tasks.

    Parameters

    ----------

    source_model: dc.models.TorchModel

        Source model to copy parameter values from.

    include_top: bool, default True

        if true, copies the last dense layer

    """

    assignment_map: Dict[Any, Any] = {}

    source_vars = list(source_model.model.parameters())

    dest_vars = list(self.model.parameters())

    if not include_top:

      source_vars = source_vars[:-2]

      dest_vars = dest_vars[:-2]

    for source_var, dest_var in zip(source_vars, dest_vars):

      assignment_map[source_var] = dest_var

    return assignment_map

  def _create_value_map(self, source_model: "TorchModel",

                        **kwargs) -> Dict[Any, Any]:

    """

    Creates a value map between parameters in the source model and their

    current values. This is used only when a custom value map is missing, and

    assumes the restore method has been called.

    Parameters

    ----------

    source_model: dc.models.TorchModel

        Source model to create value map from

    """

    value_map: Dict[Any, Any] = {}

    source_vars = list(source_model.model.parameters())

    for source_var in source_vars:

      value_map[source_var] = source_var.detach().numpy()

    return value_map

  def load_from_pretrained(self,

                           source_model: "TorchModel",

                           assignment_map: Optional[Dict[Any, Any]] = None,

                           value_map: Optional[Dict[Any, Any]] = None,

                           checkpoint: Optional[str] = None,

                           model_dir: Optional[str] = None,

                           include_top: bool = True,

                           inputs: Optional[Sequence[Any]] = None,

                           **kwargs) -> None:

    """Copies parameter values from a pretrained model. `source_model` can either

    be a pretrained model or a model with the same architecture. `value_map`

    is a parameter-value dictionary. If no `value_map` is provided, the parameter

    values are restored to the `source_model` from a checkpoint and a default

    `value_map` is created. `assignment_map` is a dictionary mapping parameters

    from the `source_model` to the current model. If no `assignment_map` is

    provided, one is made from scratch and assumes the model is composed of

    several different layers, with the final one being a dense layer. include_top

    is used to control whether or not the final dense layer is used. The default

    assignment map is useful in cases where the type of task is different

    (classification vs regression) and/or number of tasks in the setting.

    Parameters

    ----------

    source_model: dc.TorchModel, required

      source_model can either be the pretrained model or a dc.TorchModel with

      the same architecture as the pretrained model. It is used to restore from

      a checkpoint, if value_map is None and to create a default assignment map

      if assignment_map is None

    assignment_map: Dict, default None

      Dictionary mapping the source_model parameters and current model parameters

    value_map: Dict, default None

      Dictionary containing source_model trainable parameters mapped to numpy

      arrays. If value_map is None, the values are restored and a default

      parameter map is created using the restored values

    checkpoint: str, default None

      the path to the checkpoint file to load.  If this is None, the most recent

      checkpoint will be chosen automatically.  Call get_checkpoints() to get a

      list of all available checkpoints

    model_dir: str, default None

      Restore model from custom model directory if needed

    include_top: bool, default True

        if True, copies the weights and bias associated with the final dense

        layer. Used only when assignment map is None

    inputs: List, input tensors for model

        if not None, then the weights are built for both the source and self.

    """

    if inputs is not None:

      # Ensure weights for both models are built.

      source_model.model(inputs)

      self.model(inputs)

    self._ensure_built()

    if value_map is None:

      logger.info(

          "No value map provided. Creating default value map from restored model."

      )

      source_model.restore(model_dir=model_dir, checkpoint=checkpoint)

      value_map = self._create_value_map(source_model=source_model)

    if assignment_map is None:

      logger.info("No assignment map provided. Creating custom assignment map.")

      assignment_map = self._create_assignment_map(

          source_model=source_model, include_top=include_top)

    for source_var, dest_var in assignment_map.items():

      assert source_var.shape == dest_var.shape

      dest_var.data = torch.as_tensor(value_map[source_var])

class _StandardLoss(object):