class_docstr: A trainer for `AbstractComponent`s. Used mainly for supervised learning.
  
  Create a `ModelTrainer` by passing a `AbstractNeuralModel` in the constructor.
  Invoke `train()` to initiate the training loop. The root `TNeuralModule` should return a scalar loss.
class_docstr_tok: ["A","trainer","for","`","AbstractComponent","`","s",".","Used","mainly","for","supervised","learning",".","Create","a","`","ModelTrainer","`","by","passing","a","`","AbstractNeuralModel","`","in","the","constructor",".","Invoke","`","train","(",")","`","to","initiate","the","training","loop",".","The","root","`","TNeuralModule","`","should","return","a","scalar","loss","."]
docstr: :param model: The Component to be built and trained
  :param checkpoint_location: The location where the trained model will be checkpointed and saved.
  :param max_num_epochs: The maximum number of epochs to run training for.
  :param minibatch_size: The maximum size of the minibatch (`BaseComponent`s can override this
      by detecting full minibatches and returning False in `extend_minibatch_by_sample`)
  :param optimizer_creator: An optional function that accepts an iterable of the training parameters
      (pyTorch tensors) and returns a PyTorch optimizer.
  :param scheduler_creator: An optional function that accepts an optimizer and creates a scheduler
      implementing `AbstractScheduler`. This could be a wrapper for existing learning schedulers.
      The scheduler will be invoked at after each training step.
  :param clip_gradient_norm: An optional norm for clipping the gradient norms during training.
  :param target_validation_metric: An optional string of the name of the metric (returned by
      the TNeuralModule) which is used to detect if the model performance improved in validation.
      This is used for early stopping, and checkpointing the best model. If `None` the model
      loss (value returned from `forward()` of TNeuralModule) is used.
  :param target_validation_metric_higher_is_better: if `True` increases to `target_validation_metric`
      imply improvements. Ignored if `target_validation_metric` is `None`.
  :param enable_amp: Enable automatic mixed precision during training.
  :param catch_cuda_ooms: Catch CUDA out-of-memory errors (OOM) and resume training when they happen.
docstr_tok: [":","param","model",":","The","Component","to","be","built","and","trained",":","param","checkpoint_location",":","The","location","where","the","trained","model","will","be","checkpointed","and","saved",".",":","param","max_num_epochs",":","The","maximum","number","of","epochs","to","run","training","for",".",":","param","minibatch_size",":","The","maximum","size","of","the","minibatch","(","`","BaseComponent","`","s","can","override","this","by","detecting","full","minibatches","and","returning","False","in","`","extend_minibatch_by_sample","`",")",":","param","optimizer_creator",":","An","optional","function","that","accepts","an","iterable","of","the","training","parameters","(","pyTorch","tensors",")","and","returns","a","PyTorch","optimizer",".",":","param","scheduler_creator",":","An","optional","function","that","accepts","an","optimizer","and","creates","a","scheduler","implementing","`","AbstractScheduler","`",".","This","could","be","a","wrapper","for","existing","learning","schedulers",".","The","scheduler","will","be","invoked","at","after","each","training","step",".",":","param","clip_gradient_norm",":","An","optional","norm","for","clipping","the","gradient","norms","during","training",".",":","param","target_validation_metric",":","An","optional","string","of","the","name","of","the","metric","(","returned","by","the","TNeuralModule",")","which","is","used","to","detect","if","the","model","performance","improved","in","validation",".","This","is","used","for","early","stopping",",","and","checkpointing","the","best","model",".","If","`","None","`","the","model","loss","(","value","returned","from","`","forward","(",")","`","of","TNeuralModule",")","is","used",".",":","param","target_validation_metric_higher_is_better",":","if","`","True","`","increases","to","`","target_validation_metric","`","imply","improvements",".","Ignored","if","`","target_validation_metric","`","is","`","None","`",".",":","param","enable_amp",":","Enable","automatic","mixed","precision","during","training",".",":","param","catch_cuda_ooms",":","Catch","CUDA","out-of-memory","errors","(","OOM",")","and","resume","training","when","they","happen","."]
code: def __init__(
      self,
      model: AbstractNeuralModel[
          TRawDatapoint, TTensorizedDatapoint, TNeuralModule
      ],
      checkpoint_location: Path,
      *,
      max_num_epochs: int = 100,
      minibatch_size: int = 200,
      optimizer_creator: Optional[
          Callable[[Iterable[torch.Tensor]], torch.optim.Optimizer]
      ] = None,
      scheduler_creator: Optional[
          Callable[[torch.optim.Optimizer], AbstractScheduler]
      ] = None,
      clip_gradient_norm: Optional[float] = None,
      target_validation_metric: Optional[str] = None,
      target_validation_metric_higher_is_better: bool = False,
      enable_amp: bool = False,
      catch_cuda_ooms: bool = False,
  ):
      """
      :param model: The Component to be built and trained
      :param checkpoint_location: The location where the trained model will be checkpointed and saved.
      :param max_num_epochs: The maximum number of epochs to run training for.
      :param minibatch_size: The maximum size of the minibatch (`BaseComponent`s can override this
          by detecting full minibatches and returning False in `extend_minibatch_by_sample`)
      :param optimizer_creator: An optional function that accepts an iterable of the training parameters
          (pyTorch tensors) and returns a PyTorch optimizer.
      :param scheduler_creator: An optional function that accepts an optimizer and creates a scheduler
          implementing `AbstractScheduler`. This could be a wrapper for existing learning schedulers.
          The scheduler will be invoked at after each training step.
      :param clip_gradient_norm: An optional norm for clipping the gradient norms during training.
      :param target_validation_metric: An optional string of the name of the metric (returned by
          the TNeuralModule) which is used to detect if the model performance improved in validation.
          This is used for early stopping, and checkpointing the best model. If `None` the model
          loss (value returned from `forward()` of TNeuralModule) is used.
      :param target_validation_metric_higher_is_better: if `True` increases to `target_validation_metric`
          imply improvements. Ignored if `target_validation_metric` is `None`.
      :param enable_amp: Enable automatic mixed precision during training.
      :param catch_cuda_ooms: Catch CUDA out-of-memory errors (OOM) and resume training when they happen.
      """
      self.__model = model
      self.__neural_network: Optional[TNeuralModule] = None
      self._checkpoint_location = checkpoint_location
  
      self._max_num_epochs = max_num_epochs
      self._minibatch_size = minibatch_size
      if optimizer_creator is None:
          self._create_optimizer = lambda p: torch.optim.Adam(p)
      else:
          self._create_optimizer = optimizer_creator
  
      self._create_scheduler = scheduler_creator
  
      self.__metadata_finalized_hooks: List[
          Callable[[ModelType], None]
      ] = []
      self._training_start_hooks: List[
          Callable[
              [ModelType, TNeuralModule, torch.optim.Optimizer], None
          ]
      ] = []
      self._train_epoch_end_hooks: List[EndOfEpochHook] = []
      self._validation_epoch_end_hooks: List[EndOfEpochHook] = []
      self._improved_epoch_end_hooks: List[EndOfEpochHook] = []
      self._clip_gradient_norm = clip_gradient_norm
      self._enable_amp = enable_amp
      self._catch_cuda_ooms = catch_cuda_ooms
  
      self._target_metric = target_validation_metric
      if target_validation_metric is not None:
          self._target_metric_higher_is_better = (
              target_validation_metric_higher_is_better
          )
      else:
          assert (
              not target_validation_metric_higher_is_better
          ), "When no explicit metric is passed, the validation loss will be used."
          self._target_metric_higher_is_better = False
  
imports: from typing_extensions import Final, Protocol
  import json
  import logging
  import math
  import time
  import torch
  from pathlib import Path
  from tqdm import tqdm
  from typing import Callable, Dict, Generic, Iterable, List, Optional, TypeVar
  from ptgnn.baseneuralmodel.abstractneuralmodel import AbstractNeuralModel
  from ptgnn.baseneuralmodel.modulewithmetrics import ModuleWithMetrics
  from ptgnn.baseneuralmodel.utils.data import MemorizedDataIterable
  from ptgnn.baseneuralmodel.utils.oom import catch_cuda_oom
label_desc: Use image node_id 1 to create a new ModelTrainer object from inherited base classes: Generic with example: obj = ModelTrainer(model, checkpoint_location)