class_docstr: ModularTorchModel is a subclass of TorchModel that allows for components to be
  pretrained and then combined into a final model. It is designed to be subclassed
  for specific models and is not intended to be used directly. There are 3 main differences
  between ModularTorchModel and TorchModel:
  
  - The build_components() method is used to define the components of the model.
  - The components are combined into a final model with the build_model() method.
  - The loss function is defined with the loss_func method. This may access the
    components to compute the loss using intermediate values from the network, rather
    than just the full forward pass output.
  
  Here is an example of how to use ModularTorchModel to pretrain a linear layer, load
  it into another network and then finetune that network:
  
  >>> import numpy as np
  >>> import deepchem as dc
  >>> import torch
  >>> n_samples = 6
  >>> n_feat = 3
  >>> n_hidden = 2
  >>> n_tasks = 6
  >>> pt_tasks = 3
  >>> X = np.random.rand(n_samples, n_feat)
  >>> y_pretrain = np.zeros((n_samples, pt_tasks)).astype(np.float32)
  >>> dataset_pt = dc.data.NumpyDataset(X, y_pretrain)
  >>> y_finetune = np.zeros((n_samples, n_tasks)).astype(np.float32)
  >>> dataset_ft = dc.data.NumpyDataset(X, y_finetune)
  >>> components = {'linear': torch.nn.Linear(n_feat, n_hidden),
  ... 'activation': torch.nn.ReLU(), 'head': torch.nn.Linear(n_hidden, n_tasks)}
  >>> model = torch.nn.Sequential(components['linear'], components['activation'],
  ... components['head'])
  >>> modular_model = dc.models.torch_models.modular.ModularTorchModel(model, components)
  >>> def example_loss_func(inputs, labels, weights):
  ...    return (torch.nn.functional.mse_loss(model(inputs), labels[0]) * weights[0]).mean()
  >>> modular_model.loss_func = example_loss_func
  >>> def example_model_build():
  ...     return torch.nn.Sequential(components['linear'], components['activation'],
  ... components['head'])
  >>> modular_model.build_model = example_model_build
  >>> pretrain_components = {'linear': torch.nn.Linear(n_feat, n_hidden),
  ... 'activation': torch.nn.ReLU(), 'head': torch.nn.Linear(n_hidden, pt_tasks)}
  >>> pretrain_model = torch.nn.Sequential(pretrain_components['linear'],
  ... pretrain_components['activation'], pretrain_components['head'])
  >>> pretrain_modular_model = dc.models.torch_models.modular.ModularTorchModel(pretrain_model,
  ... pretrain_components)
  >>> def example_pt_loss_func(inputs, labels, weights):
  ...     return (torch.nn.functional.mse_loss(pretrain_model(inputs), labels[0]) * weights[0]).mean()
  >>> pretrain_modular_model.loss_func = example_pt_loss_func
  >>> pt_loss = pretrain_modular_model.fit(dataset_pt, nb_epoch=1)
  >>> modular_model.load_from_pretrained(pretrain_modular_model, components=['linear'])
  >>> ft_loss = modular_model.fit(dataset_ft, nb_epoch=1)
class_docstr_tok: ["ModularTorchModel","is","a","subclass","of","TorchModel","that","allows","for","components","to","be","pretrained","and","then","combined","into","a","final","model",".","It","is","designed","to","be","subclassed","for","specific","models","and","is","not","intended","to","be","used","directly",".","There","are","3","main","differences","between","ModularTorchModel","and","TorchModel",":","-","The","build_components","(",")","method","is","used","to","define","the","components","of","the","model",".","-","The","components","are","combined","into","a","final","model","with","the","build_model","(",")","method",".","-","The","loss","function","is","defined","with","the","loss_func","method",".","This","may","access","the","components","to","compute","the","loss","using","intermediate","values","from","the","network",",","rather","than","just","the","full","forward","pass","output",".","Here","is","an","example","of","how","to","use","ModularTorchModel","to","pretrain","a","linear","layer",",","load","it","into","another","network","and","then","finetune","that","network",":",">",">",">","import","numpy","as","np",">",">",">","import","deepchem","as","dc",">",">",">","import","torch",">",">",">","n_samples","=","6",">",">",">","n_feat","=","3",">",">",">","n_hidden","=","2",">",">",">","n_tasks","=","6",">",">",">","pt_tasks","=","3",">",">",">","X","=","np.random.rand","(","n_samples",",","n_feat",")",">",">",">","y_pretrain","=","np.zeros","(","(","n_samples",",","pt_tasks",")",")",".astype","(","np.float32",")",">",">",">","dataset_pt","=","dc.data.NumpyDataset","(","X",",","y_pretrain",")",">",">",">","y_finetune","=","np.zeros","(","(","n_samples",",","n_tasks",")",")",".astype","(","np.float32",")",">",">",">","dataset_ft","=","dc.data.NumpyDataset","(","X",",","y_finetune",")",">",">",">","components","=","{","'linear","'",":","torch.nn.Linear","(","n_feat",",","n_hidden",")",",","...","'activation","'",":","torch.nn.ReLU","(",")",",","'head","'",":","torch.nn.Linear","(","n_hidden",",","n_tasks",")","}",">",">",">","model","=","torch.nn.Sequential","(","components","[","'linear","'","]",",","components","[","'activation","'","]",",","...","components","[","'head","'","]",")",">",">",">","modular_model","=","dc.models.torch_models.modular.ModularTorchModel","(","model",",","components",")",">",">",">","def","example_loss_func","(","inputs",",","labels",",","weights",")",":","...","return","(","torch.nn.functional.mse_loss","(","model","(","inputs",")",",","labels","[","0","]",")","*","weights","[","0","]",")",".mean","(",")",">",">",">","modular_model.loss_func","=","example_loss_func",">",">",">","def","example_model_build","(",")",":","...","return","torch.nn.Sequential","(","components","[","'linear","'","]",",","components","[","'activation","'","]",",","...","components","[","'head","'","]",")",">",">",">","modular_model.build_model","=","example_model_build",">",">",">","pretrain_components","=","{","'linear","'",":","torch.nn.Linear","(","n_feat",",","n_hidden",")",",","...","'activation","'",":","torch.nn.ReLU","(",")",",","'head","'",":","torch.nn.Linear","(","n_hidden",",","pt_tasks",")","}",">",">",">","pretrain_model","=","torch.nn.Sequential","(","pretrain_components","[","'linear","'","]",",","...","pretrain_components","[","'activation","'","]",",","pretrain_components","[","'head","'","]",")",">",">",">","pretrain_modular_model","=","dc.models.torch_models.modular.ModularTorchModel","(","pretrain_model",",","...","pretrain_components",")",">",">",">","def","example_pt_loss_func","(","inputs",",","labels",",","weights",")",":","...","return","(","torch.nn.functional.mse_loss","(","pretrain_model","(","inputs",")",",","labels","[","0","]",")","*","weights","[","0","]",")",".mean","(",")",">",">",">","pretrain_modular_model.loss_func","=","example_pt_loss_func",">",">",">","pt_loss","=","pretrain_modular_model.fit","(","dataset_pt",",","nb_epoch=1",")",">",">",">","modular_model.load_from_pretrained","(","pretrain_modular_model",",","components=","[","'linear","'","]",")",">",">",">","ft_loss","=","modular_model.fit","(","dataset_ft",",","nb_epoch=1",")"]
docstr: Train this model on data from a generator. This method is similar to
  the TorchModel implementation, but it passes the inputs directly to the
  loss function, rather than passing them through the model first.  This
  enables the loss to be calculated from intermediate steps of the model
  and not just the final output.
  
  Parameters
  ----------
  generator: generator
      this should generate batches, each represented as a tuple of the form
      (inputs, labels, weights).
  max_checkpoints_to_keep: int
      the maximum number of checkpoints to keep.  Older checkpoints are discarded.
  checkpoint_interval: int
      the frequency at which to write checkpoints, measured in training steps.
      Set this to 0 to disable automatic checkpointing.
  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 torch.nn.Parameter
      the variables to train.  If None (the default), all trainable variables in
      the model are used.
  loss: function
      a function of the form f(outputs, labels, weights) that computes the loss
      for each batch.  If None (the default), the model's standard loss function
      is used.
  callbacks: function or list of functions
      one or more functions of the form f(model, step) that will be invoked after
      every step.  This can be used to perform validation, logging, etc.
  all_losses: Optional[List[float]], optional (default None)
      If specified, all logged losses are appended into this list. Note that
      you can call `fit()` repeatedly with the same list and losses will
      continue to be appended.
  
  Returns
  -------
  The average loss over the most recent checkpoint interval
docstr_tok: ["Train","this","model","on","data","from","a","generator",".","This","method","is","similar","to","the","TorchModel","implementation",",","but","it","passes","the","inputs","directly","to","the","loss","function",",","rather","than","passing","them","through","the","model","first",".","This","enables","the","loss","to","be","calculated","from","intermediate","steps","of","the","model","and","not","just","the","final","output",".","Parameters","--","--","--","--","--","generator",":","generator","this","should","generate","batches",",","each","represented","as","a","tuple","of","the","form","(","inputs",",","labels",",","weights",")",".","max_checkpoints_to_keep",":","int","the","maximum","number","of","checkpoints","to","keep",".","Older","checkpoints","are","discarded",".","checkpoint_interval",":","int","the","frequency","at","which","to","write","checkpoints",",","measured","in","training","steps",".","Set","this","to","0","to","disable","automatic","checkpointing",".","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","torch.nn.Parameter","the","variables","to","train",".","If","None","(","the","default",")",",","all","trainable","variables","in","the","model","are","used",".","loss",":","function","a","function","of","the","form","f","(","outputs",",","labels",",","weights",")","that","computes","the","loss","for","each","batch",".","If","None","(","the","default",")",",","the","model","'s","standard","loss","function","is","used",".","callbacks",":","function","or","list","of","functions","one","or","more","functions","of","the","form","f","(","model",",","step",")","that","will","be","invoked","after","every","step",".","This","can","be","used","to","perform","validation",",","logging",",","etc",".","all_losses",":","Optional","[","List","[","float","]","]",",","optional","(","default","None",")","If","specified",",","all","logged","losses","are","appended","into","this","list",".","Note","that","you","can","call","`","fit","(",")","`","repeatedly","with","the","same","list","and","losses","will","continue","to","be","appended",".","Returns","--","--","--","-","The","average","loss","over","the","most","recent","checkpoint","interval"]
code: def fit_generator(
      self,
      generator: Iterable[Tuple[Any, Any, Any]],
      max_checkpoints_to_keep: int = 5,
      checkpoint_interval: int = 1000,
      restore: bool = False,
      variables: Optional[List[torch.nn.Parameter]] = None,
      loss: Optional[LossFn] = None,
      callbacks: Union[Callable, List[Callable]] = [],
      all_losses: Optional[List[float]] = None,
  ) -> float:
      """Train this model on data from a generator. This method is similar to
      the TorchModel implementation, but it passes the inputs directly to the
      loss function, rather than passing them through the model first.  This
      enables the loss to be calculated from intermediate steps of the model
      and not just the final output.
  
      Parameters
      ----------
      generator: generator
          this should generate batches, each represented as a tuple of the form
          (inputs, labels, weights).
      max_checkpoints_to_keep: int
          the maximum number of checkpoints to keep.  Older checkpoints are discarded.
      checkpoint_interval: int
          the frequency at which to write checkpoints, measured in training steps.
          Set this to 0 to disable automatic checkpointing.
      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 torch.nn.Parameter
          the variables to train.  If None (the default), all trainable variables in
          the model are used.
      loss: function
          a function of the form f(outputs, labels, weights) that computes the loss
          for each batch.  If None (the default), the model's standard loss function
          is used.
      callbacks: function or list of functions
          one or more functions of the form f(model, step) that will be invoked after
          every step.  This can be used to perform validation, logging, etc.
      all_losses: Optional[List[float]], optional (default None)
          If specified, all logged losses are appended into this list. Note that
          you can call `fit()` repeatedly with the same list and losses will
          continue to be appended.
  
      Returns
      -------
      The average loss over the most recent checkpoint interval
      """
  
      if not isinstance(callbacks, SequenceCollection):
          callbacks = [callbacks]
      self._ensure_built()
      self.model.train()
      avg_loss = 0.0
      last_avg_loss = 0.0
      averaged_batches = 0
      # FIXME This line is not needed as loss is computed inside the call to loss_func
      if loss is None:
          loss = self._loss_fn
      if variables is None:
          optimizer = self._pytorch_optimizer
          lr_schedule = self._lr_schedule
      else:
          var_key = tuple(variables)
          if var_key in self._optimizer_for_vars:
              optimizer, lr_schedule = self._optimizer_for_vars[var_key]
          else:
              optimizer = self.optimizer._create_pytorch_optimizer(
                  variables
              )
              if isinstance(
                  self.optimizer.learning_rate, LearningRateSchedule
              ):
                  lr_schedule = self.optimizer.learning_rate._create_pytorch_schedule(
                      optimizer
                  )
              else:
                  lr_schedule = None
              self._optimizer_for_vars[var_key] = (
                  optimizer,
                  lr_schedule,
              )
      time1 = time.time()
  
      # Main training loop.
  
      for batch in generator:
          if restore:
              self.restore()
              restore = False
          inputs: OneOrMany[torch.Tensor]
          inputs, labels, weights = self._prepare_batch(batch)
  
          # Execute the loss function, accumulating the gradients.
  
          if isinstance(inputs, list) and len(inputs) == 1:
              inputs = inputs[0]
  
          optimizer.zero_grad()
          batch_loss = self.loss_func(inputs, labels, weights)
          batch_loss.backward()
          optimizer.step()
          if lr_schedule is not None:
              lr_schedule.step()
          self._global_step += 1
          current_step = self._global_step
  
          avg_loss += float(batch_loss)
  
          # Report progress and write checkpoints.
          averaged_batches += 1
          should_log = current_step % self.log_frequency == 0
          if should_log:
              avg_loss = float(avg_loss) / averaged_batches
              logger.info(
                  "Ending global_step %d: Average loss %.10f"
                  % (current_step, avg_loss)
              )
              if all_losses is not None:
                  all_losses.append(avg_loss)
              # Capture the last avg_loss in case of return since we're resetting to 0 now
              last_avg_loss = avg_loss
              avg_loss = 0.0
              averaged_batches = 0
          if (
              checkpoint_interval > 0
              and current_step % checkpoint_interval
              == checkpoint_interval - 1
          ):
              self.save_checkpoint(max_checkpoints_to_keep)
          for c in callbacks:
              c(self, current_step)
          if self.tensorboard and should_log:
              self._log_scalar_to_tensorboard(
                  "loss", batch_loss, current_step
              )
          if (self.wandb_logger is not None) and should_log:
              all_data = dict({"train/loss": batch_loss})
              self.wandb_logger.log_data(all_data, step=current_step)
  
      # Report final results.
      if averaged_batches > 0:
          avg_loss = float(avg_loss) / averaged_batches
          logger.info(
              "Ending global_step %d: Average loss %g"
              % (current_step, avg_loss)
          )
          if all_losses is not None:
              all_losses.append(avg_loss)
          last_avg_loss = avg_loss
  
      if checkpoint_interval > 0:
          self.save_checkpoint(max_checkpoints_to_keep)
  
      time2 = time.time()
      logger.info(
          "TIMING: model fitting took %0.3f s" % (time2 - time1)
      )
      return last_avg_loss
  
imports: import time
  import logging
  import os
  from collections.abc import Sequence
  from typing import Any, Callable, Iterable, List, Optional, Tuple, Union, Sequence
  import torch
  import torch.nn
  from deepchem.models.torch_models.torch_model import TorchModel
  from deepchem.models.optimizers import LearningRateSchedule
  from deepchem.utils.typing import LossFn, OneOrMany
label_desc: Use image node_id 7 for calling the ModularTorchModel obj's underlying member method code with example usage: obj.fit_generator(generator, max_checkpoints_to_keep, checkpoint_interval, restore, variables, loss, callbacks, all_losses) and returns: last_avg_loss