<commit_before># coding=utf-8
# Copyright 2020 The Trax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Trax main training functions."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import itertools
import os
import pickle
import random
import sys
import time

from absl import logging

import gin

import jax
import numpy
import six
import tensorflow.compat.v2 as tf
from trax import history as trax_history
from trax import jaxboard
from trax import layers as tl
from trax import lr_schedules as lr
from trax import math
from trax import optimizers as trax_opt
from trax.math import numpy as np
from trax.math import random as jax_random
from trax.shapes import ShapeDtype
from trax.supervised import inputs as trax_inputs


# TODO(afrozm): Maybe flatten everything from OptState into TrainerState.
TrainerState = collections.namedtuple('_TrainerState', [
    'step',         # Current training step number.
    'opt_state',    # OptState.
    'history',      # trax.history.History.
    'model_state',  # Auxilliary state of the model.
])


OptState = collections.namedtuple('_OptState', [
    'weights',     # Model weights.
    'slots',       # Per-parameter optimizer state, e.g. gradient moments.
    'opt_params',  # Optimizer (hyper)parameters, e.g. learning rate, momentum.
])


_DEFAULT_METRICS = {
    'loss': tl.CrossEntropyLoss,
    'accuracy': tl.AccuracyScalar,
    'sequence_accuracy': tl.SequenceAccuracyScalar,
    'neg_log_perplexity': tl.CrossEntropyLoss,
    'weights_per_batch_per_core': tl.SumOfWeights,
}


class Trainer(object):
  """Trax trainer.

  A trainer allows to make training steps, train for full epochs,
  save the training state and access evaluation data.
  """

  def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
               output_dir=None, random_seed=None, n_devices=None,
               checkpoints_at=None, should_save_checkpoints=True,
               should_write_summaries=True, has_weights=False,
               nontrainable_param_map=None, id_to_mask=None,
               metrics=None, checkpoint_highest=None, checkpoint_lowest=None):

    self._is_chief, self._n_devices, rng = (
        self._init_host_and_devices(n_devices, random_seed))
    self._should_save_checkpoints = should_save_checkpoints and self._is_chief
    self._checkpoints_at = checkpoints_at or []
    self._should_write_summaries = should_write_summaries
    if not output_dir:
      self._should_save_checkpoints = False
      self._should_write_summaries = False
    self._checkpoint_highest = checkpoint_highest
    self._checkpoint_lowest = checkpoint_lowest
    self._has_weights = has_weights
    self._id_to_mask = id_to_mask
    self._metrics_dict = metrics if metrics is not None else _DEFAULT_METRICS
    loss_fn = loss_fn(has_weights=has_weights, id_to_mask=id_to_mask)
    # Inputs is either an Inputs instance or a function that returns it.
    self._inputs = inputs
    if callable(inputs):  # If we pass a function, e.g., through gin, call it.
      self._inputs = inputs()

    # Initialize the learning rate to a dummy value. It will be set in reset().
    opt = optimizer(learning_rate=0.0)

    # Setup the model.
    model_train = model(mode='train')
    model_predict_eval = model(mode='eval')

    # Setup state.
    rng, init_rng = jax_random.split(rng)
    self._rngs = np.stack(jax_random.split(rng, self._n_devices))
    # If the inputs are a tuple/list, add [None] (batch) to each element.
    if self._inputs.input_shape and isinstance(
        self._inputs.input_shape[0], (list, tuple)
    ):
      model_input_shape = tuple(
          tuple([None] + list(shape)) for shape in self._inputs.input_shape)
    else:  # Otherwise just add [None] to the input shape.
      model_input_shape = tuple([None] + list(self._inputs.input_shape))
    # Same for targets.
    if self._inputs.target_shape and isinstance(
        self._inputs.target_shape[0], (list, tuple)
    ):
      model_target_shape = tuple(
          tuple([None] + list(shape)) for shape in self._inputs.target_shape)
    else:
      model_target_shape = tuple([None] + list(self._inputs.target_shape))
    # Change all None to 1 in input and target shape.
    model_input_shape = math.nested_map(lambda x: x or 1, model_input_shape)
    model_target_shape = math.nested_map(lambda x: x or 1,
                                         model_target_shape)

    def new_opt_state_and_model_state(shape_dtype, rng):
      """Returns optimizer and model states suitable for training a model."""
      # Combine inputs and targets on the stack.
      shapes, dtypes = shape_dtype
      input_signature = tuple(ShapeDtype(s, d)
                              for (s, d) in zip(shapes, dtypes))
      # We need to create a new model instance and not reuse `model_train` here,
      # because `m.initialize` puts cached parameter values in `m` and hence the
      # next call of `m.initialize` will give wrong results.
      m = tl.Serial(model(mode='train'), loss_fn)
      m._set_rng_recursive(rng)  # pylint: disable=protected-access
      weights, state = m.init(input_signature)
      (slots, opt_params) = opt.tree_init(weights)
      return (OptState(weights, slots, opt_params), state)

    if _is_jit_init():
      # JIT parameter initialization to avoid memory fragmentation
      new_opt_state_and_model_state = math.jit(new_opt_state_and_model_state,
                                               static_argnums=(0,))
    self._new_opt_state_and_model_state = (
        lambda: new_opt_state_and_model_state(  # pylint: disable=g-long-lambda
            self._inputs.example_shape_dtype, init_rng))

    # Arrange and initialize metrics layers.
    self._metrics = list(sorted(self._metrics_dict.keys()))
    metrics_layers = [self._metrics_dict[m](has_weights=self._has_weights,
                                            id_to_mask=self._id_to_mask)
                      for m in self._metrics]
    metrics_in_parallel = tl.Branch(*metrics_layers)
    # TODO(lukaszkaiser): clean this up once layer API stabilizes.
    # For now, we need to initialize metric layers somehow, so here we go.
    # We assume that they do not have any parameters, so this is a dummy.
    dummy_shapes = ((1, 2), (1,), (1,)) if self._has_weights else ((1, 2), (1,))
    dummy_signature = tuple(ShapeDtype(s) for s in dummy_shapes)
    metrics_in_parallel._set_rng_recursive(init_rng)  # pylint: disable=protected-access
    m_weights, m_state = metrics_in_parallel.init(dummy_signature)
    self._metrics_weights = self._for_n_devices(m_weights)
    self._metrics_state = self._for_n_devices(m_state)

    # Jit model_predict and update so they're fast.
    self._jit_eval = _jit_predict_fn(<commit_msg>Calculate proper input signatures for metrics in Trainer.<commit_after># coding=utf-8
# Copyright 2020 The Trax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Trax main training functions."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import itertools
import os
import pickle
import random
import sys
import time

from absl import logging

import gin

import jax
import numpy
import six
import tensorflow.compat.v2 as tf
from trax import history as trax_history
from trax import jaxboard
from trax import layers as tl
from trax import lr_schedules as lr
from trax import math
from trax import optimizers as trax_opt
from trax.math import numpy as np
from trax.math import random as jax_random
from trax.shapes import ShapeDtype
from trax.supervised import inputs as trax_inputs


# TODO(afrozm): Maybe flatten everything from OptState into TrainerState.
TrainerState = collections.namedtuple('_TrainerState', [
    'step',         # Current training step number.
    'opt_state',    # OptState.
    'history',      # trax.history.History.
    'model_state',  # Auxilliary state of the model.
])


OptState = collections.namedtuple('_OptState', [
    'weights',     # Model weights.
    'slots',       # Per-parameter optimizer state, e.g. gradient moments.
    'opt_params',  # Optimizer (hyper)parameters, e.g. learning rate, momentum.
])


_DEFAULT_METRICS = {
    'loss': tl.CrossEntropyLoss,
    'accuracy': tl.AccuracyScalar,
    'sequence_accuracy': tl.SequenceAccuracyScalar,
    'neg_log_perplexity': tl.CrossEntropyLoss,
    'weights_per_batch_per_core': tl.SumOfWeights,
}


class Trainer(object):
  """Trax trainer.

  A trainer allows to make training steps, train for full epochs,
  save the training state and access evaluation data.
  """

  def __init__(self, model, loss_fn, optimizer, lr_schedule, inputs,
               output_dir=None, random_seed=None, n_devices=None,
               checkpoints_at=None, should_save_checkpoints=True,
               should_write_summaries=True, has_weights=False,
               nontrainable_param_map=None, id_to_mask=None,
               metrics=None, checkpoint_highest=None, checkpoint_lowest=None):

    self._is_chief, self._n_devices, rng = (
        self._init_host_and_devices(n_devices, random_seed))
    self._should_save_checkpoints = should_save_checkpoints and self._is_chief
    self._checkpoints_at = checkpoints_at or []
    self._should_write_summaries = should_write_summaries
    if not output_dir:
      self._should_save_checkpoints = False
      self._should_write_summaries = False
    self._checkpoint_highest = checkpoint_highest
    self._checkpoint_lowest = checkpoint_lowest
    self._has_weights = has_weights
    self._id_to_mask = id_to_mask
    self._metrics_dict = metrics if metrics is not None else _DEFAULT_METRICS
    loss_fn = loss_fn(has_weights=has_weights, id_to_mask=id_to_mask)
    # Inputs is either an Inputs instance or a function that returns it.
    self._inputs = inputs
    if callable(inputs):  # If we pass a function, e.g., through gin, call it.
      self._inputs = inputs()

    # Initialize the learning rate to a dummy value. It will be set in reset().
    opt = optimizer(learning_rate=0.0)

    # Setup the model.
    model_train = model(mode='train')
    model_predict_eval = model(mode='eval')

    # Setup state.
    rng, init_rng = jax_random.split(rng)
    self._rngs = np.stack(jax_random.split(rng, self._n_devices))
    # If the inputs are a tuple/list, add [None] (batch) to each element.
    if self._inputs.input_shape and isinstance(
        self._inputs.input_shape[0], (list, tuple)
    ):
      model_input_shape = tuple(
          tuple([None] + list(shape)) for shape in self._inputs.input_shape)
    else:  # Otherwise just add [None] to the input shape.
      model_input_shape = tuple([None] + list(self._inputs.input_shape))
    # Same for targets.
    if self._inputs.target_shape and isinstance(
        self._inputs.target_shape[0], (list, tuple)
    ):
      model_target_shape = tuple(
          tuple([None] + list(shape)) for shape in self._inputs.target_shape)
    else:
      model_target_shape = tuple([None] + list(self._inputs.target_shape))
    # Change all None to 1 in input and target shape.
    model_input_shape = math.nested_map(lambda x: x or 1, model_input_shape)
    model_target_shape = math.nested_map(lambda x: x or 1,
                                         model_target_shape)

    def new_opt_state_and_model_state(shape_dtype, rng):
      """Returns optimizer and model states suitable for training a model."""
      # Combine inputs and targets on the stack.
      shapes, dtypes = shape_dtype
      input_signature = tuple(ShapeDtype(s, d)
                              for (s, d) in zip(shapes, dtypes))
      # We need to create a new model instance and not reuse `model_train` here,
      # because `m.initialize` puts cached parameter values in `m` and hence the
      # next call of `m.initialize` will give wrong results.
      m = tl.Serial(model(mode='train'), loss_fn)
      m._set_rng_recursive(rng)  # pylint: disable=protected-access
      weights, state = m.init(input_signature)
      (slots, opt_params) = opt.tree_init(weights)
      return (OptState(weights, slots, opt_params), state)

    if _is_jit_init():
      # JIT parameter initialization to avoid memory fragmentation
      new_opt_state_and_model_state = math.jit(new_opt_state_and_model_state,
                                               static_argnums=(0,))
    self._new_opt_state_and_model_state = (
        lambda: new_opt_state_and_model_state(  # pylint: disable=g-long-lambda
            self._inputs.example_shape_dtype, init_rng))

    # Arrange and initialize metrics layers.
    self._metrics = list(sorted(self._metrics_dict.keys()))
    metrics_layers = [self._metrics_dict[m](has_weights=self._has_weights,
                                            id_to_mask=self._id_to_mask)
                      for m in self._metrics]
    metrics_in_parallel = tl.Branch(*metrics_layers)
    metrics_in_parallel._set_rng_recursive(init_rng)  # pylint: disable=protected-access
    example_signature = tuple(
        ShapeDtype(s, d) for (s, d) in zip(*self._inputs.example_shape_dtype)
    )
    model_predict_eval.init(example_signature)
    output_signature = model_predict_eval.output_signature(example_signature)
    m_weights, m_state = metrics_in_parallel.init(output_signature)
    self._metrics_weights = self._for_n_devices(m_weights)
    self._metrics_state = self._for_n_devices(m_state)

    # Jit model_predict and update so they're fast.
    self._jit_eval = _jit_predict_fn(
{'fraction_non_alphanumeric': 0.08174032202809181, 'fraction_numerical': 0.0029462144570058238, 'mean_word_length': 3.419013018468059, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '643567', 'n_tokens_mistral': 4466, 'n_tokens_neox': 4371, 'n_words': 1278}