<gh_stars>10-100
from typing import Any, Dict, Generator, List, Optional

import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter

from probnmn.config import Config
from probnmn.utils.checkpointing import CheckpointManager


class _Trainer(object):
    r"""
    A base class for generic training of models. This class can have multiple models interacting
    with each other, rather than a single model, which is suitable to our use-case (for example,
    ``module_training`` phase has two models:
    :class:`~probnmn.models.program_generator.ProgramGenerator` and
    :class:`~probnmn.models.nmn.NeuralModuleNetwork`). It offers full flexibility, with sensible
    defaults which may be changed (or disabled) while extending this class.

    Extended Summary
    ----------------
    1. Default :class:`~torch.optim.Adam` Optimizer, updates parameters of all models in this
       trainer. Learning rate and weight decay for this optimizer are picked up from the provided
       config.

    2. Default :class:`~torch.optim.lr_scheduler.ReduceLROnPlateau` learning rate scheduler. Gamma
       and patience arguments are picked up from the provided config. Observed metric is assumed
       to be of type "higher is better". For 'lower is better" metrics, make sure to reciprocate.

    3. Tensorboard logging of loss curves, metrics etc.

    4. Serialization of models and optimizer as checkpoint (.pth) files after every validation.
       The observed metric for keeping track of best checkpoint is of type "higher is better",
       follow (2) above if the observed metric is of type "lower is better".

    Extend this class and override suitable methods as per requirements, some important ones are:

    1. :meth:`step`, provides complete customization, this is the method which comprises of one
       full training iteration, and internally calls (in order) - :meth:`_before_iteration`,
       :meth:`_do_iteration` and :meth:`_after_iteration`. Most of the times you may not require
       overriding this method, instead one of the mentioned three methods called by `:meth:`step`.

    2. :meth:`_do_iteration`, with core training loop - what happens every iteration, given a
       ``batch`` from the dataloader this class holds.

    3. :meth:`_before_iteration` and :meth:`_after_iteration`, for any pre- or post-processing
       steps. Default behaviour:

        * :meth:`_before_iteration` - call ``optimizer.zero_grad()``
        * :meth:`_after_iteration` - call ``optimizer.step()`` and do tensorboard logging.

    4. :meth:`after_validation`, to specify any steps after evaluation. Default behaviour is to
       do learning rate scheduling and log validation metrics on tensorboard.

    Notes
    -----
    All models are `passed by assignment`, so they could be shared with an external evaluator.
    Do not set ``self._models = ...`` anywhere while extending this class.

    Parameters
    ----------
    config: Config
        A :class:`~probnmn.Config` object with all the relevant configuration parameters.
    dataloader: torch.utils.data.DataLoader
        A :class:`~torch.utils.data.DataLoader` which provides batches of training examples. It
        wraps one of :mod:`probnmn.data.datasets` depending on the evaluation phase.
    models: Dict[str, Type[nn.Module]]
        All the models which interact with each other during training. These are one or more from
        :mod:`probnmn.models` depending on the training phase.
    serialization_dir: str
        Path to a directory for tensorboard logging and serializing checkpoints.
    gpu_ids: List[int], optional (default=[0])
        List of GPU IDs to use or evaluation, ``[-1]`` - use CPU.
    """

    def __init__(
        self,
        config: Config,
        dataloader: DataLoader,
        models: Dict[str, nn.Module],
        serialization_dir: str,
        gpu_ids: List[int] = [0],
    ):
        self._C = config

        # Make dataloader cyclic for sampling batches perpetually.
        self._dataloader = self._cycle(dataloader)
        self._models = models

        # Set device according to specified GPU ids.
        self._device = torch.device(f"cuda:{gpu_ids[0]}" if gpu_ids[0] >= 0 else "cpu")

        # Shift models to device, and wrap in DataParallel for Multi-GPU execution (if needed).
        for model_name in self._models:
            self._models[model_name] = self._models[model_name].to(self._device)

            if len(gpu_ids) > 1 and -1 not in gpu_ids:
                # Don't wrap to DataParallel if single GPU ID or -1 (CPU) is provided.
                self._models[model_name] = nn.DataParallel(self._models[model_name], gpu_ids)

        # Accumulate parameters of all models to construct Adam Optimizer.
        all_parameters: List[Any] = []
        for model_name in self._models:
            all_parameters.extend(list(self._models[model_name].parameters()))
        self._optimizer = optim.Adam(
            all_parameters, lr=self._C.OPTIM.LR_INITIAL, weight_decay=self._C.OPTIM.WEIGHT_DECAY
        )

        # Default learning rate scheduler: (lr *= gamma) when observed metric plateaus for
        # "patience" number of validation steps.
        self._lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            self._optimizer,
            mode="max",
            factor=self._C.OPTIM.LR_GAMMA,
            patience=self._C.OPTIM.LR_PATIENCE,
            threshold=1e-3,
        )

        # Tensorboard summary writer for logging losses and metrics.
        self._tensorboard_writer = SummaryWriter(log_dir=serialization_dir)

        # Checkpoint manager to serialize model, optimizer and lr scheduler periodically.
        self._checkpoint_manager = CheckpointManager(
            serialization_dir=serialization_dir,
            keep_recent=100,
            optimizer=self._optimizer,
            scheduler=self._lr_scheduler,
            **models,
        )
        # Initialize a counter to keep track of the iteration number.
        # This increments everytime ``step`` is called.
        self._iteration: int = -1

    def step(self, iteration: Optional[int] = None):
        r"""
        Perform one iteration of training.

        Parameters
        ----------
        iteration: int, optional (default = None)
            Iteration number (useful to hard set to any number when loading checkpoint).
            If ``None``, use the internal :attr:`self._iteration` counter.
        """
        self._before_iteration()

        batch = next(self._dataloader)
        output_dict = self._do_iteration(batch)
        self._after_iteration(output_dict)

        self._iteration = iteration or self._iteration + 1

    def _before_iteration(self):
        r"""
        Steps to do before doing the forward pass of iteration. Default behavior is to simply
        call :meth:`zero_grad` for optimizer. Called inside :meth:`step`.
        """
        self._optimizer.zero_grad()

    def _do_iteration(self, batch: Dict[str, Any]) -> Dict[str, Any]:
        r"""
        Forward and backward passes on models, given a batch sampled from dataloader.

        Parameters
        ----------
        batch: Dict[str, Any]
            A batch of training examples sampled from dataloader. See :func:`step` and
            :meth:`_cycle` on how this batch is sampled.

        Returns
        -------
        Dict[str, Any]
            An output dictionary typically returned by the models. This would be passed to
            :meth:`_after_iteration` for tensorboard logging.
        """
        # What a single iteration usually would look like.
        iteration_output_dict = self._models["model"](batch)
        batch_loss = iteration_output_dict["loss"].mean()
        batch_loss.backward()
        return {"loss": batch_loss}

    def _after_iteration(self, output_dict: Dict[str, Any]):
        r"""
        Steps to do after doing the forward pass of iteration. Default behavior is to simply
        do gradient update through ``optimizer.step()``, and log metrics to tensorboard.

        Parameters
        ----------
        output_dict: Dict[str, Any]
            This is exactly the object returned by :meth:_do_iteration`, which would contain all
            the required losses for tensorboard logging.
        """
        self._optimizer.step()

        # keys: {"loss"} + ... {other keys such as "elbo"}
        for key in output_dict:
            if isinstance(output_dict[key], dict):
                # Use ``add_scalars`` for dicts in a nested ``output_dict``.
                self._tensorboard_writer.add_scalars(
                    f"train/{key}", output_dict[key], self._iteration
                )
            else:
                # Use ``add_scalar`` for floats / zero-dim tensors in ``output_dict``.
                self._tensorboard_writer.add_scalar(
                    f"train/{key}", output_dict[key], self._iteration
                )

    def after_validation(self, val_metrics: Dict[str, Any], iteration: Optional[int] = None):
        r"""
        Steps to do after an external :class:`~probnmn.evaluators._evaluator._Evaluator` performs
        evaluation. This is not called by :meth:`step`, call it from outside at appropriate time.
        Default behavior is to perform learning rate scheduling, serializaing checkpoint and to
        log validation metrics to tensorboard.

        Since this implementation assumes a key ``"metric"`` in ``val_metrics``, it is convenient
        to set this key while overriding this method, when there are multiple models and multiple
        metrics and there is one metric which decides best checkpoint.

        Parameters
        ----------
        val_metrics: Dict[str, Any]
            Validation metrics for all the models. Returned by ``evaluate`` method of
            :class:`~probnmn.evaluators._evaluator._Evaluator` (or its extended class).
        iteration: int, optional (default = None)
            Iteration number. If ``None``, use the internal :attr:`self._iteration` counter.
        """
        if iteration is not None:
            self._iteration = iteration

        # Serialize model and optimizer and keep track of best checkpoint.
        self._checkpoint_manager.step(self._iteration, val_metrics["metric"])

        # Perform learning rate scheduling based on validation perplexity.
        self._lr_scheduler.step(val_metrics["metric"])

        # Log learning rate after scheduling.
        self._tensorboard_writer.add_scalar(
            "train/lr", self._optimizer.param_groups[0]["lr"], self._iteration
        )

        # Log all validation metrics to tensorboard (pop the "metric" key, which was only relevant
        # to learning rate scheduling and checkpointing).
        val_metrics.pop("metric")
        for model_name in val_metrics:
            for metric_name in val_metrics[model_name]:
                self._tensorboard_writer.add_scalar(
                    f"val/metrics/{model_name}/{metric_name}",
                    val_metrics[model_name][metric_name],
                    self._iteration,
                )

    def load_checkpoint(self, checkpoint_path: str, iteration: Optional[int] = None):
        r"""
        Load a checkpoint to continue training from. The iteration when this checkpoint was
        serialized, is inferred from its name (so do not rename after serialization).

        Parameters
        ----------
        checkpoint_path: str
            Path to a checkpoint containing models and optimizers of the phase which is being
            trained on.

        iteration: int, optional (default = None)
            Iteration number. If ``None``, get it from the checkpoint.
        """
        _iteration = self._checkpoint_manager.load(checkpoint_path)

        # By default, the provided iteration overrides what is found in checkpoint.
        iteration = iteration or _iteration
        self._iteration = iteration

    def _cycle(self, dataloader: DataLoader) -> Generator[Dict[str, torch.Tensor], None, None]:
        r"""
        A generator which yields a random batch from dataloader perpetually. This generator is
        used in the constructor.

        Extended Summary
        ----------------
        This is done so because we train for a fixed number of iterations, and do not have the
        notion of 'epochs'. Using ``itertools.cycle`` with dataloader is harmful and may cause
        unexpeced memory leaks.
        """
        while True:
            for batch in dataloader:
                for key in batch:
                    batch[key] = batch[key].to(self._device)
                yield batch

    @property
    def iteration(self):
        return self._iteration

    @property
    def models(self):
        return self._models

from typing import Any, Dict, Generator, List, Optional

import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter

from probnmn.config import Config
from probnmn.utils.checkpointing import CheckpointManager


class _Trainer(object):
    r"""
    A base class for generic training of models. This class can have multiple models interacting
    with each other, rather than a single model, which is suitable to our use-case (for example,
    ``module_training`` phase has two models:
    :class:`~probnmn.models.program_generator.ProgramGenerator` and
    :class:`~probnmn.models.nmn.NeuralModuleNetwork`). It offers full flexibility, with sensible
    defaults which may be changed (or disabled) while extending this class.

    Extended Summary
    ----------------
    1. Default :class:`~torch.optim.Adam` Optimizer, updates parameters of all models in this
       trainer. Learning rate and weight decay for this optimizer are picked up from the provided
       config.

    2. Default :class:`~torch.optim.lr_scheduler.ReduceLROnPlateau` learning rate scheduler. Gamma
       and patience arguments are picked up from the provided config. Observed metric is assumed
       to be of type "higher is better". For 'lower is better" metrics, make sure to reciprocate.

    3. Tensorboard logging of loss curves, metrics etc.

    4. Serialization of models and optimizer as checkpoint (.pth) files after every validation.
       The observed metric for keeping track of best checkpoint is of type "higher is better",
       follow (2) above if the observed metric is of type "lower is better".

    Extend this class and override suitable methods as per requirements, some important ones are:

    1. :meth:`step`, provides complete customization, this is the method which comprises of one
       full training iteration, and internally calls (in order) - :meth:`_before_iteration`,
       :meth:`_do_iteration` and :meth:`_after_iteration`. Most of the times you may not require
       overriding this method, instead one of the mentioned three methods called by `:meth:`step`.

    2. :meth:`_do_iteration`, with core training loop - what happens every iteration, given a
       ``batch`` from the dataloader this class holds.

    3. :meth:`_before_iteration` and :meth:`_after_iteration`, for any pre- or post-processing
       steps. Default behaviour:

        * :meth:`_before_iteration` - call ``optimizer.zero_grad()``
        * :meth:`_after_iteration` - call ``optimizer.step()`` and do tensorboard logging.

    4. :meth:`after_validation`, to specify any steps after evaluation. Default behaviour is to
       do learning rate scheduling and log validation metrics on tensorboard.

    Notes
    -----
    All models are `passed by assignment`, so they could be shared with an external evaluator.
    Do not set ``self._models = ...`` anywhere while extending this class.

    Parameters
    ----------
    config: Config
        A :class:`~probnmn.Config` object with all the relevant configuration parameters.
    dataloader: torch.utils.data.DataLoader
        A :class:`~torch.utils.data.DataLoader` which provides batches of training examples. It
        wraps one of :mod:`probnmn.data.datasets` depending on the evaluation phase.
    models: Dict[str, Type[nn.Module]]
        All the models which interact with each other during training. These are one or more from
        :mod:`probnmn.models` depending on the training phase.
    serialization_dir: str
        Path to a directory for tensorboard logging and serializing checkpoints.
    gpu_ids: List[int], optional (default=[0])
        List of GPU IDs to use or evaluation, ``[-1]`` - use CPU.
    """

    def __init__(
        self,
        config: Config,
        dataloader: DataLoader,
        models: Dict[str, nn.Module],
        serialization_dir: str,
        gpu_ids: List[int] = [0],
    ):
        self._C = config

        # Make dataloader cyclic for sampling batches perpetually.
        self._dataloader = self._cycle(dataloader)
        self._models = models

        # Set device according to specified GPU ids.
        self._device = torch.device(f"cuda:{gpu_ids[0]}" if gpu_ids[0] >= 0 else "cpu")

        # Shift models to device, and wrap in DataParallel for Multi-GPU execution (if needed).
        for model_name in self._models:
            self._models[model_name] = self._models[model_name].to(self._device)

            if len(gpu_ids) > 1 and -1 not in gpu_ids:
                # Don't wrap to DataParallel if single GPU ID or -1 (CPU) is provided.
                self._models[model_name] = nn.DataParallel(self._models[model_name], gpu_ids)

        # Accumulate parameters of all models to construct Adam Optimizer.
        all_parameters: List[Any] = []
        for model_name in self._models:
            all_parameters.extend(list(self._models[model_name].parameters()))
        self._optimizer = optim.Adam(
            all_parameters, lr=self._C.OPTIM.LR_INITIAL, weight_decay=self._C.OPTIM.WEIGHT_DECAY
        )

        # Default learning rate scheduler: (lr *= gamma) when observed metric plateaus for
        # "patience" number of validation steps.
        self._lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            self._optimizer,
            mode="max",
            factor=self._C.OPTIM.LR_GAMMA,
            patience=self._C.OPTIM.LR_PATIENCE,
            threshold=1e-3,
        )

        # Tensorboard summary writer for logging losses and metrics.
        self._tensorboard_writer = SummaryWriter(log_dir=serialization_dir)

        # Checkpoint manager to serialize model, optimizer and lr scheduler periodically.
        self._checkpoint_manager = CheckpointManager(
            serialization_dir=serialization_dir,
            keep_recent=100,
            optimizer=self._optimizer,
            scheduler=self._lr_scheduler,
            **models,
        )
        # Initialize a counter to keep track of the iteration number.
        # This increments everytime ``step`` is called.
        self._iteration: int = -1

    def step(self, iteration: Optional[int] = None):
        r"""
        Perform one iteration of training.

        Parameters
        ----------
        iteration: int, optional (default = None)
            Iteration number (useful to hard set to any number when loading checkpoint).
            If ``None``, use the internal :attr:`self._iteration` counter.
        """
        self._before_iteration()

        batch = next(self._dataloader)
        output_dict = self._do_iteration(batch)
        self._after_iteration(output_dict)

        self._iteration = iteration or self._iteration + 1

    def _before_iteration(self):
        r"""
        Steps to do before doing the forward pass of iteration. Default behavior is to simply
        call :meth:`zero_grad` for optimizer. Called inside :meth:`step`.
        """
        self._optimizer.zero_grad()

    def _do_iteration(self, batch: Dict[str, Any]) -> Dict[str, Any]:
        r"""
        Forward and backward passes on models, given a batch sampled from dataloader.

        Parameters
        ----------
        batch: Dict[str, Any]
            A batch of training examples sampled from dataloader. See :func:`step` and
            :meth:`_cycle` on how this batch is sampled.

        Returns
        -------
        Dict[str, Any]
            An output dictionary typically returned by the models. This would be passed to
            :meth:`_after_iteration` for tensorboard logging.
        """
        # What a single iteration usually would look like.
        iteration_output_dict = self._models["model"](batch)
        batch_loss = iteration_output_dict["loss"].mean()
        batch_loss.backward()
        return {"loss": batch_loss}

    def _after_iteration(self, output_dict: Dict[str, Any]):
        r"""
        Steps to do after doing the forward pass of iteration. Default behavior is to simply
        do gradient update through ``optimizer.step()``, and log metrics to tensorboard.

        Parameters
        ----------
        output_dict: Dict[str, Any]
            This is exactly the object returned by :meth:_do_iteration`, which would contain all
            the required losses for tensorboard logging.
        """
        self._optimizer.step()

        # keys: {"loss"} + ... {other keys such as "elbo"}
        for key in output_dict:
            if isinstance(output_dict[key], dict):
                # Use ``add_scalars`` for dicts in a nested ``output_dict``.
                self._tensorboard_writer.add_scalars(
                    f"train/{key}", output_dict[key], self._iteration
                )
            else:
                # Use ``add_scalar`` for floats / zero-dim tensors in ``output_dict``.
                self._tensorboard_writer.add_scalar(
                    f"train/{key}", output_dict[key], self._iteration
                )

    def after_validation(self, val_metrics: Dict[str, Any], iteration: Optional[int] = None):
        r"""
        Steps to do after an external :class:`~probnmn.evaluators._evaluator._Evaluator` performs
        evaluation. This is not called by :meth:`step`, call it from outside at appropriate time.
        Default behavior is to perform learning rate scheduling, serializaing checkpoint and to
        log validation metrics to tensorboard.

        Since this implementation assumes a key ``"metric"`` in ``val_metrics``, it is convenient
        to set this key while overriding this method, when there are multiple models and multiple
        metrics and there is one metric which decides best checkpoint.

        Parameters
        ----------
        val_metrics: Dict[str, Any]
            Validation metrics for all the models. Returned by ``evaluate`` method of
            :class:`~probnmn.evaluators._evaluator._Evaluator` (or its extended class).
        iteration: int, optional (default = None)
            Iteration number. If ``None``, use the internal :attr:`self._iteration` counter.
        """
        if iteration is not None:
            self._iteration = iteration

        # Serialize model and optimizer and keep track of best checkpoint.
        self._checkpoint_manager.step(self._iteration, val_metrics["metric"])

        # Perform learning rate scheduling based on validation perplexity.
        self._lr_scheduler.step(val_metrics["metric"])

        # Log learning rate after scheduling.
        self._tensorboard_writer.add_scalar(
            "train/lr", self._optimizer.param_groups[0]["lr"], self._iteration
        )

        # Log all validation metrics to tensorboard (pop the "metric" key, which was only relevant
        # to learning rate scheduling and checkpointing).
        val_metrics.pop("metric")
        for model_name in val_metrics:
            for metric_name in val_metrics[model_name]:
                self._tensorboard_writer.add_scalar(
                    f"val/metrics/{model_name}/{metric_name}",
                    val_metrics[model_name][metric_name],
                    self._iteration,
                )

    def load_checkpoint(self, checkpoint_path: str, iteration: Optional[int] = None):
        r"""
        Load a checkpoint to continue training from. The iteration when this checkpoint was
        serialized, is inferred from its name (so do not rename after serialization).

        Parameters
        ----------
        checkpoint_path: str
            Path to a checkpoint containing models and optimizers of the phase which is being
            trained on.

        iteration: int, optional (default = None)
            Iteration number. If ``None``, get it from the checkpoint.
        """
        _iteration = self._checkpoint_manager.load(checkpoint_path)

        # By default, the provided iteration overrides what is found in checkpoint.
        iteration = iteration or _iteration
        self._iteration = iteration

    def _cycle(self, dataloader: DataLoader) -> Generator[Dict[str, torch.Tensor], None, None]:
        r"""
        A generator which yields a random batch from dataloader perpetually. This generator is
        used in the constructor.

        Extended Summary
        ----------------
        This is done so because we train for a fixed number of iterations, and do not have the
        notion of 'epochs'. Using ``itertools.cycle`` with dataloader is harmful and may cause
        unexpeced memory leaks.
        """
        while True:
            for batch in dataloader:
                for key in batch:
                    batch[key] = batch[key].to(self._device)
                yield batch

    @property
    def iteration(self):
        return self._iteration

    @property
    def models(self):
        return self._models
A base class for generic training of models. This class can have multiple models interacting     with each other, rather than a single model, which is suitable to our use-case (for example,     ``module_training`` phase has two models:     :class:`~probnmn.models.program_generator.ProgramGenerator` and     :class:`~probnmn.models.nmn.NeuralModuleNetwork`). It offers full flexibility, with sensible     defaults which may be changed (or disabled) while extending this class.      Extended Summary     ----------------     1. Default :class:`~torch.optim.Adam` Optimizer, updates parameters of all models in this        trainer. Learning rate and weight decay for this optimizer are picked up from the provided        config.      2. Default :class:`~torch.optim.lr_scheduler.ReduceLROnPlateau` learning rate scheduler. Gamma        and patience arguments are picked up from the provided config. Observed metric is assumed        to be of type "higher is better". For 'lower is better" metrics, make sure to reciprocate.      3. Tensorboard logging of loss curves, metrics etc.      4. Serialization of models and optimizer as checkpoint (.pth) files after every validation.        The observed metric for keeping track of best checkpoint is of type "higher is better",        follow (2) above if the observed metric is of type "lower is better".      Extend this class and override suitable methods as per requirements, some important ones are:      1. :meth:`step`, provides complete customization, this is the method which comprises of one        full training iteration, and internally calls (in order) - :meth:`_before_iteration`,        :meth:`_do_iteration` and :meth:`_after_iteration`. Most of the times you may not require        overriding this method, instead one of the mentioned three methods called by `:meth:`step`.      2. :meth:`_do_iteration`, with core training loop - what happens every iteration, given a        ``batch`` from the dataloader this class holds.      3. :meth:`_before_iteration` and :meth:`_after_iteration`, for any pre- or post-processing        steps. Default behaviour:          * :meth:`_before_iteration` - call ``optimizer.zero_grad()``         * :meth:`_after_iteration` - call ``optimizer.step()`` and do tensorboard logging.      4. :meth:`after_validation`, to specify any steps after evaluation. Default behaviour is to        do learning rate scheduling and log validation metrics on tensorboard.      Notes     -----     All models are `passed by assignment`, so they could be shared with an external evaluator.     Do not set ``self._models = ...`` anywhere while extending this class.      Parameters     ----------     config: Config         A :class:`~probnmn.Config` object with all the relevant configuration parameters.     dataloader: torch.utils.data.DataLoader         A :class:`~torch.utils.data.DataLoader` which provides batches of training examples. It         wraps one of :mod:`probnmn.data.datasets` depending on the evaluation phase.     models: Dict[str, Type[nn.Module]]         All the models which interact with each other during training. These are one or more from         :mod:`probnmn.models` depending on the training phase.     serialization_dir: str         Path to a directory for tensorboard logging and serializing checkpoints.     gpu_ids: List[int], optional (default=[0])         List of GPU IDs to use or evaluation, ``[-1]`` - use CPU. # Make dataloader cyclic for sampling batches perpetually. # Set device according to specified GPU ids. # Shift models to device, and wrap in DataParallel for Multi-GPU execution (if needed). # Don't wrap to DataParallel if single GPU ID or -1 (CPU) is provided. # Accumulate parameters of all models to construct Adam Optimizer. # Default learning rate scheduler: (lr *= gamma) when observed metric plateaus for # "patience" number of validation steps. # Tensorboard summary writer for logging losses and metrics. # Checkpoint manager to serialize model, optimizer and lr scheduler periodically. # Initialize a counter to keep track of the iteration number. # This increments everytime ``step`` is called. Perform one iteration of training.          Parameters         ----------         iteration: int, optional (default = None)             Iteration number (useful to hard set to any number when loading checkpoint).             If ``None``, use the internal :attr:`self._iteration` counter. Steps to do before doing the forward pass of iteration. Default behavior is to simply         call :meth:`zero_grad` for optimizer. Called inside :meth:`step`. Forward and backward passes on models, given a batch sampled from dataloader.          Parameters         ----------         batch: Dict[str, Any]             A batch of training examples sampled from dataloader. See :func:`step` and             :meth:`_cycle` on how this batch is sampled.          Returns         -------         Dict[str, Any]             An output dictionary typically returned by the models. This would be passed to             :meth:`_after_iteration` for tensorboard logging. # What a single iteration usually would look like. Steps to do after doing the forward pass of iteration. Default behavior is to simply         do gradient update through ``optimizer.step()``, and log metrics to tensorboard.          Parameters         ----------         output_dict: Dict[str, Any]             This is exactly the object returned by :meth:_do_iteration`, which would contain all             the required losses for tensorboard logging. # keys: {"loss"} + ... {other keys such as "elbo"} # Use ``add_scalars`` for dicts in a nested ``output_dict``. # Use ``add_scalar`` for floats / zero-dim tensors in ``output_dict``. Steps to do after an external :class:`~probnmn.evaluators._evaluator._Evaluator` performs         evaluation. This is not called by :meth:`step`, call it from outside at appropriate time.         Default behavior is to perform learning rate scheduling, serializaing checkpoint and to         log validation metrics to tensorboard.          Since this implementation assumes a key ``"metric"`` in ``val_metrics``, it is convenient         to set this key while overriding this method, when there are multiple models and multiple         metrics and there is one metric which decides best checkpoint.          Parameters         ----------         val_metrics: Dict[str, Any]             Validation metrics for all the models. Returned by ``evaluate`` method of             :class:`~probnmn.evaluators._evaluator._Evaluator` (or its extended class).         iteration: int, optional (default = None)             Iteration number. If ``None``, use the internal :attr:`self._iteration` counter. # Serialize model and optimizer and keep track of best checkpoint. # Perform learning rate scheduling based on validation perplexity. # Log learning rate after scheduling. # Log all validation metrics to tensorboard (pop the "metric" key, which was only relevant # to learning rate scheduling and checkpointing). Load a checkpoint to continue training from. The iteration when this checkpoint was         serialized, is inferred from its name (so do not rename after serialization).          Parameters         ----------         checkpoint_path: str             Path to a checkpoint containing models and optimizers of the phase which is being             trained on.          iteration: int, optional (default = None)             Iteration number. If ``None``, get it from the checkpoint. # By default, the provided iteration overrides what is found in checkpoint. A generator which yields a random batch from dataloader perpetually. This generator is         used in the constructor.          Extended Summary         ----------------         This is done so because we train for a fixed number of iterations, and do not have the         notion of 'epochs'. Using ``itertools.cycle`` with dataloader is harmful and may cause         unexpeced memory leaks.