<commit_before>import logging
import os
import shutil
import time
from typing import Dict, Optional, List  # pylint: disable=unused-import

import torch
from torch.nn.utils.clip_grad import clip_grad_norm
import tqdm
from tensorboard import SummaryWriter

from allennlp.common import Params
from allennlp.common.checks import ConfigurationError
from allennlp.data import Dataset
from allennlp.data.iterators.data_iterator import DataIterator
from allennlp.models.model import Model
from allennlp.nn.util import arrays_to_variables, device_mapping

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name


class Trainer:
    def __init__(self,
                 model: Model,
                 optimizer: torch.optim.Optimizer,
                 iterator: DataIterator,
                 train_dataset: Dataset,
                 validation_dataset: Optional[Dataset] = None,
                 patience: int = 2,
                 validation_metric: str = "-loss",
                 num_epochs: int = 20,
                 serialization_dir: Optional[str] = None,
                 cuda_device: int = -1,
                 grad_norm: Optional[float] = None,
                 grad_clipping: Optional[float] = None,
                 no_tqdm: bool = False) -> None:
        """
        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.
        optimizer : ``torch.nn.Optimizer``, required.
            An instance of a Pytorch Optimizer, instantiated with the parameters of the
            model to be optimized.
        iterator : ``DataIterator``, required.
            A method for iterating over a ``Dataset``, yielding padded indexed batches.
        train_dataset : ``Dataset``, required.
            A ``Dataset`` to train on. The dataset should have already been indexed.
        validation_dataset : ``Dataset``, optional, (default = None).
            A ``Dataset`` to evaluate on. The dataset should have already been indexed.
        patience : int, optional (default=2)
            Number of epochs to be patient before early stopping.
        validation_metric : 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.
        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.
        cuda_device : int, optional (default = -1)
            An integer specifying the CUDA device to use. If -1, the CPU is used.
            Multi-gpu training is not currently supported, but will be once the
            Pytorch DataParallel API stabilises.
        grad_norm : float, optional, (default = None).
            If provided, gradient norms will be rescaled to have a maximum of this value.
        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.
        no_tqdm : ``bool``, optional (default=False)
            We use ``tqdm`` for logging, which will print a nice progress bar that updates in place
            after every batch.  This is nice if you're running training on a local shell, but can
            cause problems with log files from, e.g., a docker image running on kubernetes.  If
            ``no_tqdm`` is ``True``, we will not use tqdm, and instead log batch statistics using
            ``logger.info``, outputting a line at most every 10 seconds.
        """
        self._model = model
        self._iterator = iterator
        self._optimizer = optimizer
        self._train_dataset = train_dataset
        self._validation_dataset = validation_dataset

        self._patience = patience
        self._num_epochs = num_epochs
        self._serialization_dir = serialization_dir
        self._cuda_device = cuda_device
        self._grad_norm = grad_norm
        self._grad_clipping = grad_clipping

        increase_or_decrease = validation_metric[0]
        if increase_or_decrease not in ["+", "-"]:
            raise ConfigurationError("Validation metrics must specify whether they should increase "
                                     "or decrease by pre-pending the metric name with a +/-.")
        self._validation_metric = validation_metric[1:]
        self._validation_metric_decreases = increase_or_decrease == "-"
        self._no_tqdm = no_tqdm

        if self._cuda_device >= 0:
            self._model = self._model.cuda(self._cuda_device)

        self._log_interval = 10  # seconds

    def train(self) -> None:
        epoch_counter = 0
        # Resume from serialization path if it contains a saved model.
        if self._serialization_dir is not None:
            # Set up tensorboard logging.
            train_log = SummaryWriter(os.path.join(self._serialization_dir, "log", "train"))
            validation_log = SummaryWriter(os.path.join(self._serialization_dir, "log", "validation"))
            if any(["model_state_epoch_" in x
                    for x in os.listdir(self._serialization_dir)]):
                logger.info("Loading model from checkpoint.")
                epoch_counter = self._restore_checkpoint()

        if self._grad_clipping is not None:
            # Pylint is unable to tell that we're in the case that _glad_clipping is not None...
            # pylint: disable=invalid-unary-operand-type
            clip_function = lambda grad: grad.clamp(-self._grad_clipping, self._grad_clipping)
            for parameter in self._model.parameters():
                if parameter.requires_grad:
                    parameter.register_hook(clip_function)

        logger.info("Beginning training.")
        num_training_batches = self._iterator.get_num_batches(self._train_dataset)
        if self._validation_dataset is not None:
            num_validation_batches = self._iterator.get_num_batches(self._validation_dataset)
        validation_metric_per_epoch = []  # type: List[float]
        for epoch in range(epoch_counter, self._num_epochs):
            logger.info("Epoch %d/%d", epoch + 1, self._num_epochs)
            train_loss = 0.0
            val_loss = 0.0
            # Set the model to "train" mode.
            self._model.train()
            train_generator = self._iterator(self._train_dataset, num_epochs=1)

            train_generator_tqdm = tqdm.tqdm(train_generator,
                                             disable=self._no_tqdm,
                                             total=num_training_batches)
            last_log = time.time()
            batch_num = 0
            logger.info("Training")
            for batch in train_generator_tqdm:
                batch_num += 1
                tensor_batch = arrays_to_variables(batch, self._cuda_device)
                self._optimizer.zero_grad()
                output_dict = self._model.forward(**tensor_batch)
                try:
                    loss = output_dict["loss"]
                    loss.backward()
                    # Make sure Variable is on the cpu before converting to numpy.
                    # .cpu() is a no-op if you aren't using GPUs.
                    train_loss += loss.data.cpu().numpy()
                except KeyError:
                    raise ConfigurationError("The model you are trying to optimize does not contain a"
                                             " 'loss' key in the output of model.forward(inputs).")

                if self._grad_norm:
                    clip_grad_norm(self._model.parameters(), self._grad_norm)
                self._optimizer.step()
                metrics = self._model.get_metrics()
                metrics["loss"] = float(train_loss / batch_num)
                description = self._description_from_metrics(metrics)
                train_generator_tqdm.set_description(description)
                if self._no_tqdm and time.time() - last_log > self._log_interval:
                    logger.info("Batch %d/%d: %s", batch_num, num_training_batches, description)
                    last_log = time.time()
            metrics = self._model.get_metrics(reset=True)
            metrics["loss"] = float(train_loss / batch_num)

            if self._validation_dataset is not None:
                logger.info("Validating")
                # Switch to evaluation mode.
                self._model.eval()
                val_generator = self._iterator(self._validation_dataset, num_epochs=1)
                val_generator_tqdm = tqdm.tqdm(val_generator,
                                               disable=self._no_tqdm,
                                               total=num_validation_batches)
                batch_num = 0
                for batch in val_generator_tqdm:
                    batch_num += 1
                    tensor_batch = arrays_to_variables(batch, self._cuda_device, for_training=False)
                    val_output_dict = self._model.forward(**tensor_batch)
                    loss = val_output_dict["loss"]
                    val_loss += loss.data.cpu().numpy()<commit_msg>Track parameters in tensorboard (#196)<commit_after>import logging
import os
import shutil
import time
from typing import Dict, Optional, List  # pylint: disable=unused-import

import torch
from torch.nn.utils.clip_grad import clip_grad_norm
import tqdm
from tensorboard import SummaryWriter

from allennlp.common import Params
from allennlp.common.checks import ConfigurationError
from allennlp.data import Dataset
from allennlp.data.iterators.data_iterator import DataIterator
from allennlp.models.model import Model
from allennlp.nn.util import arrays_to_variables, device_mapping

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name


class Trainer:
    def __init__(self,
                 model: Model,
                 optimizer: torch.optim.Optimizer,
                 iterator: DataIterator,
                 train_dataset: Dataset,
                 validation_dataset: Optional[Dataset] = None,
                 patience: int = 2,
                 validation_metric: str = "-loss",
                 num_epochs: int = 20,
                 serialization_dir: Optional[str] = None,
                 cuda_device: int = -1,
                 grad_norm: Optional[float] = None,
                 grad_clipping: Optional[float] = None,
                 no_tqdm: bool = False) -> None:
        """
        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.
        optimizer : ``torch.nn.Optimizer``, required.
            An instance of a Pytorch Optimizer, instantiated with the parameters of the
            model to be optimized.
        iterator : ``DataIterator``, required.
            A method for iterating over a ``Dataset``, yielding padded indexed batches.
        train_dataset : ``Dataset``, required.
            A ``Dataset`` to train on. The dataset should have already been indexed.
        validation_dataset : ``Dataset``, optional, (default = None).
            A ``Dataset`` to evaluate on. The dataset should have already been indexed.
        patience : int, optional (default=2)
            Number of epochs to be patient before early stopping.
        validation_metric : 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.
        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.
        cuda_device : int, optional (default = -1)
            An integer specifying the CUDA device to use. If -1, the CPU is used.
            Multi-gpu training is not currently supported, but will be once the
            Pytorch DataParallel API stabilises.
        grad_norm : float, optional, (default = None).
            If provided, gradient norms will be rescaled to have a maximum of this value.
        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.
        no_tqdm : ``bool``, optional (default=False)
            We use ``tqdm`` for logging, which will print a nice progress bar that updates in place
            after every batch.  This is nice if you're running training on a local shell, but can
            cause problems with log files from, e.g., a docker image running on kubernetes.  If
            ``no_tqdm`` is ``True``, we will not use tqdm, and instead log batch statistics using
            ``logger.info``, outputting a line at most every 10 seconds.
        """
        self._model = model
        self._iterator = iterator
        self._optimizer = optimizer
        self._train_dataset = train_dataset
        self._validation_dataset = validation_dataset

        self._patience = patience
        self._num_epochs = num_epochs
        self._serialization_dir = serialization_dir
        self._cuda_device = cuda_device
        self._grad_norm = grad_norm
        self._grad_clipping = grad_clipping

        increase_or_decrease = validation_metric[0]
        if increase_or_decrease not in ["+", "-"]:
            raise ConfigurationError("Validation metrics must specify whether they should increase "
                                     "or decrease by pre-pending the metric name with a +/-.")
        self._validation_metric = validation_metric[1:]
        self._validation_metric_decreases = increase_or_decrease == "-"
        self._no_tqdm = no_tqdm

        if self._cuda_device >= 0:
            self._model = self._model.cuda(self._cuda_device)

        self._log_interval = 10  # seconds
        self._summary_interval = 100  # num batches between logging to tensorboard

    def train(self) -> None:
        epoch_counter = 0
        # Resume from serialization path if it contains a saved model.
        if self._serialization_dir is not None:
            # Set up tensorboard logging.
            train_log = SummaryWriter(os.path.join(self._serialization_dir, "log", "train"))
            validation_log = SummaryWriter(os.path.join(self._serialization_dir, "log", "validation"))
            if any(["model_state_epoch_" in x
                    for x in os.listdir(self._serialization_dir)]):
                logger.info("Loading model from checkpoint.")
                epoch_counter = self._restore_checkpoint()

        if self._grad_clipping is not None:
            # Pylint is unable to tell that we're in the case that _glad_clipping is not None...
            # pylint: disable=invalid-unary-operand-type
            clip_function = lambda grad: grad.clamp(-self._grad_clipping, self._grad_clipping)
            for parameter in self._model.parameters():
                if parameter.requires_grad:
                    parameter.register_hook(clip_function)

        logger.info("Beginning training.")
        num_training_batches = self._iterator.get_num_batches(self._train_dataset)
        if self._validation_dataset is not None:
            num_validation_batches = self._iterator.get_num_batches(self._validation_dataset)
        validation_metric_per_epoch = []  # type: List[float]
        for epoch in range(epoch_counter, self._num_epochs):
            logger.info("Epoch %d/%d", epoch + 1, self._num_epochs)
            train_loss = 0.0
            val_loss = 0.0
            # Set the model to "train" mode.
            self._model.train()
            train_generator = self._iterator(self._train_dataset, num_epochs=1)

            train_generator_tqdm = tqdm.tqdm(train_generator,
                                             disable=self._no_tqdm,
                                             total=num_training_batches)
            last_log = time.time()
            batch_num = 0
            logger.info("Training")
            for batch in train_generator_tqdm:
                batch_num += 1
                tensor_batch = arrays_to_variables(batch, self._cuda_device)
                self._optimizer.zero_grad()
                output_dict = self._model.forward(**tensor_batch)
                try:
                    loss = output_dict["loss"]
                    loss.backward()
                    # Make sure Variable is on the cpu before converting to numpy.
                    # .cpu() is a no-op if you aren't using GPUs.
                    train_loss += loss.data.cpu().numpy()
                except KeyError:
                    raise ConfigurationError("The model you are trying to optimize does not contain a"
                                             " 'loss' key in the output of model.forward(inputs).")

                if self._grad_norm:
                    clip_grad_norm(self._model.parameters(), self._grad_norm)
                self._optimizer.step()
                metrics = self._model.get_metrics()
                metrics["loss"] = float(train_loss / batch_num)
                description = self._description_from_metrics(metrics)
                train_generator_tqdm.set_description(description)

                batch_num_total = num_training_batches * epoch + batch_num
                if self._serialization_dir and batch_num_total % self._summary_interval == 0:
                    for name, param in self._model.named_parameters():
                        train_log.add_scalar("PARAMETER_MEAN/" + name, param.data.mean(), batch_num_total)
                        train_log.add_scalar("PARAMETER_STD/" + name, param.data.std(), batch_num_total)
                        if param.grad is not None:
                            train_log.add_scalar("GRAD_MEAN/" + name, param.grad.data.mean(), batch_num_total)
                            train_log.add_scalar("GRAD_STD/" + name, param.grad.data.std(), batch_num_total)
                    train_log.add_scalar("LOSS/loss_train", metrics["loss"], batch_num_total)
                if self._no_tqdm and time.time() - last_log > self._log_interval:
                    logger.info("Batch %d/%d: %s", batch_num, num_training_batches, description)
                    last_log = time.time()
            metrics = self._model.get_metrics(reset=True)
            metrics["loss"] = float(train_loss / batch_num)

            if self._validation_dataset is not None:
                logger.info("Validating")
                # Switch to evaluation mode.
                self._model.eval()
                val_generator = self._iterator(self._validation_dataset, num_epochs=1)
                val_generator_tqdm = tqdm.tqdm(val_generator,
                                               disable=self._no_tqdm,
                                               total=num_validation_batches)
                batch_num = 0
                for batch in val_generator_tqdm:
                    batch_num += 1
                    tensor_batch = arrays_to_variables(batch, self._cuda_device, for_training=False)
                    val_output_dict = self._model.forward(**tensor_batch)
                    loss = val_output_dict["loss"]
                    val_loss += loss.data.cpu().numpy()
{'fraction_non_alphanumeric': 0.0769118363370677, 'fraction_numerical': 0.0030686799805163174, 'mean_word_length': 1.994166545136357, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 24, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3074016', 'n_tokens_mistral': 5334, 'n_tokens_neox': 5134, 'n_words': 1598}