class_docstr: A trainer for doing supervised learning with gradient descent. It just takes a labeled dataset
  and a `DataLoader`, and uses the supplied `Optimizer` to learn the weights for your model over
  some fixed number of epochs. You can also pass in a validation data_loader and enable early
  stopping. There are many other bells and whistles as well.
  
  Registered as a `Trainer` with the name "gradient_descent" (and is also the default `Trainer`).
  The constructor that is registered is [`from_partial_objects`](#from_partial_objects) -
  see the arguments to that function for the exact keys that should be used, if you are using
  a configuration file. They largely match the arguments to `__init__`, and we don't repeat their
  docstrings in `from_partial_objects`.
  
  [0]: https://tinyurl.com/y5mv44fw
  
  # Parameters
  
  model : `Model`, required.
      An AllenNLP model to be optimized. Pytorch Modules can also be optimized if
      their `forward` method returns a dictionary with a "loss" key, containing a
      scalar tensor representing the loss function to be optimized.
  
      If you are training your model using GPUs, your model should already be
      on the correct device. (If you are using our `train` command this will be
      handled for you.)
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  optimizer : `torch.nn.Optimizer`, required.
      An instance of a Pytorch Optimizer, instantiated with the parameters of the
      model to be optimized.
  
  data_loader : `DataLoader`, required.
      A `DataLoader` containing your `Dataset`, yielding padded indexed batches.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  patience : `Optional[int] > 0`, optional (default=`None`)
      Number of epochs to be patient before early stopping: the training is stopped
      after `patience` epochs with no improvement. If given, it must be `> 0`.
      If None, early stopping is disabled.
  
  validation_metric : `Union[str, List[str]]`, optional (default=`"-loss"`)
      Validation metric to measure for whether to stop training using patience
      and whether to serialize an `is_best` model each epoch. The metric name
      must be prepended with either "+" or "-", which specifies whether the metric
      is an increasing or decreasing function. If you specify more than one metric,
      the metrics will be summed to make the `is_best` decision.
  
  validation_data_loader : `DataLoader`, optional (default=`None`)
      A `DataLoader` to use for the validation set.  If `None`, then
      use the training `DataLoader` with the validation data.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  num_epochs : `int`, optional (default = `20`)
      Number of training epochs.
  
  serialization_dir : `str`, optional (default=`None`)
      Path to directory for saving and loading model files. Models will not be saved if
      this parameter is not passed.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  checkpointer : `Checkpointer`, optional (default=`None`)
      A `Checkpointer` is responsible for periodically saving model weights.  If none is given
      here, we will construct one with default parameters.
  
  cuda_device : `Optional[Union[int, torch.device]]`, optional (default = `None`)
      An integer or `torch.device` specifying the CUDA device to use for this process.
      If -1, the CPU is used. If `None` and you have a GPU available, that GPU will be used.
  
      !!! Note
          If you *don't* intend to use a GPU, but you have one available, you'll need
          to explicitly set `cuda_device=-1`.
  
      !!! Note
          If you intend to use a GPU, your model already needs to be on the correct device,
          which you can do with `model = model.cuda()`.
  
      !!! Note
          Data parallelism is controlled at the allennlp train level, so each trainer will have a single GPU.
  
  grad_norm : `Union[float, bool]`, optional (default = `False`)
      If a float, gradient norms will be rescaled to have a maximum of this value.
      If `True`, the gradient norms will be calculated and passed through to any `TrainerCallbacks`,
      but won't be rescaled.
      If `False`, gradient norms will not be calculated or rescaled.
  
  grad_clipping : `float`, optional (default = `None`)
      If provided, gradients will be clipped `during the backward pass` to have an (absolute)
      maximum of this value.  If you are getting `NaNs` in your gradients during training
      that are not solved by using `grad_norm`, you may need this.
  
  learning_rate_scheduler : `LearningRateScheduler`, optional (default = `None`)
      If specified, the learning rate will be decayed with respect to
      this schedule at the end of each epoch (or batch, if the scheduler implements
      the `step_batch` method). If you use `torch.optim.lr_scheduler.ReduceLROnPlateau`,
      this will use the `validation_metric` provided to determine if learning has plateaued.
      To support updating the learning rate on every batch, this can optionally implement
      `step_batch(batch_num_total)` which updates the learning rate given the batch number.
  
  momentum_scheduler : `MomentumScheduler`, optional (default = `None`)
      If specified, the momentum will be updated at the end of each batch or epoch
      according to the schedule.
  
  moving_average : `MovingAverage`, optional, (default = `None`)
      If provided, we will maintain moving averages for all parameters. During training, we
      employ a shadow variable for each parameter, which maintains the moving average. During
      evaluation, we backup the original parameters and assign the moving averages to corresponding
      parameters. Be careful that when saving the checkpoint, we will save the moving averages of
      parameters. This is necessary because we want the saved model to perform as well as the validated
      model if we load it later. But this may cause problems if you restart the training from checkpoint.
  
  callbacks : `List[TrainerCallback]`, optional (default = `None`)
      A list of callbacks that can be called at certain events: e.g. each batch, epoch, and at the start
      and end of training, etc.
  
  distributed : `bool`, optional, (default = `False`)
      If set, PyTorch's `DistributedDataParallel` is used to train the model in multiple GPUs. This also
      requires `world_size` to be greater than 1.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately (you need a top-level "distributed" key, next to
      the "trainer" entry, that specifies a list of "cuda_devices").
  
  local_rank : `int`, optional, (default = `0`)
      This is the unique identifier of the `Trainer` in a distributed process group. The GPU device id is
      used as the rank.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  world_size : `int`, (default = `1`)
      The number of `Trainer` workers participating in the distributed training.
  
      In a typical AllenNLP configuration file, this parameter does not get an entry under the
      "trainer", it gets constructed separately.
  
  num_gradient_accumulation_steps : `int`, optional, (default = `1`)
      Gradients are accumulated for the given number of steps before doing an optimizer step. This can
      be useful to accommodate batches that are larger than the RAM size. Refer [Thomas Wolf's
      post][0] for details on Gradient Accumulation.
  
  use_amp : `bool`, optional, (default = `False`)
      If `True`, we'll train using [Automatic Mixed Precision](https://pytorch.org/docs/stable/amp.html).
  
  enable_default_callbacks : `bool`, optional (default = `True`)
      When `True`, the [`DEFAULT_CALLBACKS`](#default_callbacks) will be used in
      addition to any other callbacks listed in the `callbacks` parameter.
      When set to `False`, `DEFAULT_CALLBACKS` are not used.
  
  run_confidence_checks : `bool`, optional (default = `True`)
      Determines whether model confidence checks, such as
      [`NormalizationBiasVerification`](../../confidence_checks/normalization_bias_verification/),
      are run.
  
  run_sanity_checks : `bool`, optional (default = `True`)
      This parameter is deprecated. Please use `run_confidence_checks` instead.
  
  grad_scaling : `bool`, optional (default = `True`)
      When `use_amp` is `True`, this determines whether or not to use a [`GradScaler`]
      (https://pytorch.org/docs/stable/amp.html?highlight=gradscaler#torch.cuda.amp.GradScaler).
  
      !!! Note
          This parameter is ignored when `use_amp` is `False`.
  
  ddp_wrapped_model : `Optional[DdpWrappedModel]`, optional (default = `None`)
      The `model` wrapped with a `DdpAccelerator` for distributed training.
  
      !!! Note
          This is required for distributed training.
class_docstr_tok: ["A","trainer","for","doing","supervised","learning","with","gradient","descent",".","It","just","takes","a","labeled","dataset","and","a","`","DataLoader","`",",","and","uses","the","supplied","`","Optimizer","`","to","learn","the","weights","for","your","model","over","some","fixed","number","of","epochs",".","You","can","also","pass","in","a","validation","data_loader","and","enable","early","stopping",".","There","are","many","other","bells","and","whistles","as","well",".","Registered","as","a","`","Trainer","`","with","the","name","``","gradient_descent","''","(","and","is","also","the","default","`","Trainer","`",")",".","The","constructor","that","is","registered","is","[","`","from_partial_objects","`","]","(","#","from_partial_objects",")","-","see","the","arguments","to","that","function","for","the","exact","keys","that","should","be","used",",","if","you","are","using","a","configuration","file",".","They","largely","match","the","arguments","to","`","__init__","`",",","and","we","do","n't","repeat","their","docstrings","in","`","from_partial_objects","`",".","[","0","]",":","https",":","\/\/tinyurl.com\/y5mv44fw","#","Parameters","model",":","`","Model","`",",","required",".","An","AllenNLP","model","to","be","optimized",".","Pytorch","Modules","can","also","be","optimized","if","their","`","forward","`","method","returns","a","dictionary","with","a","``","loss","''","key",",","containing","a","scalar","tensor","representing","the","loss","function","to","be","optimized",".","If","you","are","training","your","model","using","GPUs",",","your","model","should","already","be","on","the","correct","device",".","(","If","you","are","using","our","`","train","`","command","this","will","be","handled","for","you",".",")","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","optimizer",":","`","torch.nn.Optimizer","`",",","required",".","An","instance","of","a","Pytorch","Optimizer",",","instantiated","with","the","parameters","of","the","model","to","be","optimized",".","data_loader",":","`","DataLoader","`",",","required",".","A","`","DataLoader","`","containing","your","`","Dataset","`",",","yielding","padded","indexed","batches",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","patience",":","`","Optional","[","int","]",">","0","`",",","optional","(","default=","`","None","`",")","Number","of","epochs","to","be","patient","before","early","stopping",":","the","training","is","stopped","after","`","patience","`","epochs","with","no","improvement",".","If","given",",","it","must","be","`",">","0","`",".","If","None",",","early","stopping","is","disabled",".","validation_metric",":","`","Union","[","str",",","List","[","str","]","]","`",",","optional","(","default=","`","``","-loss","''","`",")","Validation","metric","to","measure","for","whether","to","stop","training","using","patience","and","whether","to","serialize","an","`","is_best","`","model","each","epoch",".","The","metric","name","must","be","prepended","with","either","``","+","''","or","``","-","''",",","which","specifies","whether","the","metric","is","an","increasing","or","decreasing","function",".","If","you","specify","more","than","one","metric",",","the","metrics","will","be","summed","to","make","the","`","is_best","`","decision",".","validation_data_loader",":","`","DataLoader","`",",","optional","(","default=","`","None","`",")","A","`","DataLoader","`","to","use","for","the","validation","set",".","If","`","None","`",",","then","use","the","training","`","DataLoader","`","with","the","validation","data",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","num_epochs",":","`","int","`",",","optional","(","default","=","`","20","`",")","Number","of","training","epochs",".","serialization_dir",":","`","str","`",",","optional","(","default=","`","None","`",")","Path","to","directory","for","saving","and","loading","model","files",".","Models","will","not","be","saved","if","this","parameter","is","not","passed",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","checkpointer",":","`","Checkpointer","`",",","optional","(","default=","`","None","`",")","A","`","Checkpointer","`","is","responsible","for","periodically","saving","model","weights",".","If","none","is","given","here",",","we","will","construct","one","with","default","parameters",".","cuda_device",":","`","Optional","[","Union","[","int",",","torch.device","]","]","`",",","optional","(","default","=","`","None","`",")","An","integer","or","`","torch.device","`","specifying","the","CUDA","device","to","use","for","this","process",".","If","-1",",","the","CPU","is","used",".","If","`","None","`","and","you","have","a","GPU","available",",","that","GPU","will","be","used",".","!","!","!","Note","If","you","*","do","n't","*","intend","to","use","a","GPU",",","but","you","have","one","available",",","you","'ll","need","to","explicitly","set","`","cuda_device=-1","`",".","!","!","!","Note","If","you","intend","to","use","a","GPU",",","your","model","already","needs","to","be","on","the","correct","device",",","which","you","can","do","with","`","model","=","model.cuda","(",")","`",".","!","!","!","Note","Data","parallelism","is","controlled","at","the","allennlp","train","level",",","so","each","trainer","will","have","a","single","GPU",".","grad_norm",":","`","Union","[","float",",","bool","]","`",",","optional","(","default","=","`","False","`",")","If","a","float",",","gradient","norms","will","be","rescaled","to","have","a","maximum","of","this","value",".","If","`","True","`",",","the","gradient","norms","will","be","calculated","and","passed","through","to","any","`","TrainerCallbacks","`",",","but","wo","n't","be","rescaled",".","If","`","False","`",",","gradient","norms","will","not","be","calculated","or","rescaled",".","grad_clipping",":","`","float","`",",","optional","(","default","=","`","None","`",")","If","provided",",","gradients","will","be","clipped","`","during","the","backward","pass","`","to","have","an","(","absolute",")","maximum","of","this","value",".","If","you","are","getting","`","NaNs","`","in","your","gradients","during","training","that","are","not","solved","by","using","`","grad_norm","`",",","you","may","need","this",".","learning_rate_scheduler",":","`","LearningRateScheduler","`",",","optional","(","default","=","`","None","`",")","If","specified",",","the","learning","rate","will","be","decayed","with","respect","to","this","schedule","at","the","end","of","each","epoch","(","or","batch",",","if","the","scheduler","implements","the","`","step_batch","`","method",")",".","If","you","use","`","torch.optim.lr_scheduler.ReduceLROnPlateau","`",",","this","will","use","the","`","validation_metric","`","provided","to","determine","if","learning","has","plateaued",".","To","support","updating","the","learning","rate","on","every","batch",",","this","can","optionally","implement","`","step_batch","(","batch_num_total",")","`","which","updates","the","learning","rate","given","the","batch","number",".","momentum_scheduler",":","`","MomentumScheduler","`",",","optional","(","default","=","`","None","`",")","If","specified",",","the","momentum","will","be","updated","at","the","end","of","each","batch","or","epoch","according","to","the","schedule",".","moving_average",":","`","MovingAverage","`",",","optional",",","(","default","=","`","None","`",")","If","provided",",","we","will","maintain","moving","averages","for","all","parameters",".","During","training",",","we","employ","a","shadow","variable","for","each","parameter",",","which","maintains","the","moving","average",".","During","evaluation",",","we","backup","the","original","parameters","and","assign","the","moving","averages","to","corresponding","parameters",".","Be","careful","that","when","saving","the","checkpoint",",","we","will","save","the","moving","averages","of","parameters",".","This","is","necessary","because","we","want","the","saved","model","to","perform","as","well","as","the","validated","model","if","we","load","it","later",".","But","this","may","cause","problems","if","you","restart","the","training","from","checkpoint",".","callbacks",":","`","List","[","TrainerCallback","]","`",",","optional","(","default","=","`","None","`",")","A","list","of","callbacks","that","can","be","called","at","certain","events",":","e.g",".","each","batch",",","epoch",",","and","at","the","start","and","end","of","training",",","etc",".","distributed",":","`","bool","`",",","optional",",","(","default","=","`","False","`",")","If","set",",","PyTorch","'s","`","DistributedDataParallel","`","is","used","to","train","the","model","in","multiple","GPUs",".","This","also","requires","`","world_size","`","to","be","greater","than","1",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately","(","you","need","a","top-level","``","distributed","''","key",",","next","to","the","``","trainer","''","entry",",","that","specifies","a","list","of","``","cuda_devices","''",")",".","local_rank",":","`","int","`",",","optional",",","(","default","=","`","0","`",")","This","is","the","unique","identifier","of","the","`","Trainer","`","in","a","distributed","process","group",".","The","GPU","device","id","is","used","as","the","rank",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","world_size",":","`","int","`",",","(","default","=","`","1","`",")","The","number","of","`","Trainer","`","workers","participating","in","the","distributed","training",".","In","a","typical","AllenNLP","configuration","file",",","this","parameter","does","not","get","an","entry","under","the","``","trainer","''",",","it","gets","constructed","separately",".","num_gradient_accumulation_steps",":","`","int","`",",","optional",",","(","default","=","`","1","`",")","Gradients","are","accumulated","for","the","given","number","of","steps","before","doing","an","optimizer","step",".","This","can","be","useful","to","accommodate","batches","that","are","larger","than","the","RAM","size",".","Refer","[","Thomas","Wolf's","post","]","[","0","]","for","details","on","Gradient","Accumulation",".","use_amp",":","`","bool","`",",","optional",",","(","default","=","`","False","`",")","If","`","True","`",",","we","'ll","train","using","[","Automatic","Mixed","Precision","]","(","https",":","\/\/pytorch.org\/docs\/stable\/amp.html",")",".","enable_default_callbacks",":","`","bool","`",",","optional","(","default","=","`","True","`",")","When","`","True","`",",","the","[","`","DEFAULT_CALLBACKS","`","]","(","#","default_callbacks",")","will","be","used","in","addition","to","any","other","callbacks","listed","in","the","`","callbacks","`","parameter",".","When","set","to","`","False","`",",","`","DEFAULT_CALLBACKS","`","are","not","used",".","run_confidence_checks",":","`","bool","`",",","optional","(","default","=","`","True","`",")","Determines","whether","model","confidence","checks",",","such","as","[","`","NormalizationBiasVerification","`","]","(","..","\/","..","\/confidence_checks\/normalization_bias_verification\/",")",",","are","run",".","run_sanity_checks",":","`","bool","`",",","optional","(","default","=","`","True","`",")","This","parameter","is","deprecated",".","Please","use","`","run_confidence_checks","`","instead",".","grad_scaling",":","`","bool","`",",","optional","(","default","=","`","True","`",")","When","`","use_amp","`","is","`","True","`",",","this","determines","whether","or","not","to","use","a","[","`","GradScaler","`","]","(","https",":","\/\/pytorch.org\/docs\/stable\/amp.html","?","highlight=gradscaler","#","torch.cuda.amp.GradScaler",")",".","!","!","!","Note","This","parameter","is","ignored","when","`","use_amp","`","is","`","False","`",".","ddp_wrapped_model",":","`","Optional","[","DdpWrappedModel","]","`",",","optional","(","default","=","`","None","`",")","The","`","model","`","wrapped","with","a","`","DdpAccelerator","`","for","distributed","training",".","!","!","!","Note","This","is","required","for","distributed","training","."]
code: def _load_model_state(self, path: str) -> None:
      # This function is only called after training. So load model on the CPU.
      device = torch.device("cpu")
      if self._ddp_wrapped_model is not None:
          self._ddp_wrapped_model.load_state_dict(
              torch.load(path, map_location=device)
          )
      else:
          self._pytorch_model.load_state_dict(
              torch.load(path, map_location=device)
          )
  
imports: import datetime
  import glob
  import logging
  import math
  import os
  import re
  import time
  import warnings
  from typing import Optional, Union, List, Dict, Tuple, Any, Type
  import torch
  from torch.cuda import amp
  from torch.nn.utils import clip_grad_norm_
  import torch.distributed
  from torch.cuda.amp.grad_scaler import OptState
  from allennlp.common.checks import ConfigurationError, check_for_gpu
  from allennlp.common import util, Tqdm, Lazy
  from allennlp.common.file_utils import hardlink_or_copy
  from allennlp.data.data_loaders.data_loader import DataLoader, TensorDict
  from allennlp.models.model import Model
  from allennlp.nn.parallel import DdpAccelerator, DdpWrappedModel, TorchDdpAccelerator
  from allennlp.nn.util import dist_reduce_sum
  from allennlp.training.callbacks import ConsoleLoggerCallback
  from allennlp.training.callbacks.confidence_checks import ConfidenceChecksCallback
  from allennlp.training.callbacks.backward import MixedPrecisionBackwardCallback
  from allennlp.training.checkpointer import Checkpointer
  from allennlp.training.learning_rate_schedulers.learning_rate_scheduler import LearningRateScheduler
  from allennlp.training.metric_tracker import MetricTracker
  from allennlp.training.momentum_schedulers.momentum_scheduler import MomentumScheduler
  from allennlp.training.moving_average import MovingAverage
  from allennlp.training.optimizers import Optimizer
  from allennlp.training.trainer import Trainer, TrainerCheckpoint
  from allennlp.training.callbacks import TrainerCallback
  from allennlp.training import util
label_desc: Use image node_id 11 for calling the GradientDescentTrainer obj's underlying member method code with example usage: obj._load_model_state(path) without return types