train_test.py

# Copyright 2020 The Flax 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.

# Lint as: python3
"""Tests for flax.examples.seq2seq.train."""

import functools

from absl.testing import absltest
import jax
from jax import random
import numpy as np

from flax import nn
import train

jax.config.parse_flags_with_absl()


class TrainTest(absltest.TestCase):

  def test_character_table(self):
    text = '410+19'
    enc_text = train.CTABLE.encode(text)
    dec_text = train.CTABLE.decode(enc_text)
    # The text is possibly padded with whitespace, but the trimmed output should
    # be equal to the input.
    self.assertEqual(text, dec_text.strip())

  def test_onehot(self):
    np.testing.assert_equal(
        train.onehot(np.array([0, 1, 2]), 4),
        np.array(
            [[1, 0, 0, 0],
             [0, 1, 0, 0],
             [0, 0, 1, 0]],
            dtype=np.float32)
    )
    np.testing.assert_equal(
        jax.vmap(functools.partial(train.onehot, vocab_size=4))(
            np.array([[0, 1], [2, 3]])),
        np.array(
            [[[1, 0, 0, 0],
              [0, 1, 0, 0]],
             [[0, 0, 1, 0],
              [0, 0, 0, 1]]],
            dtype=np.float32)
    )

  def test_get_sequence_lengths(self):
    oh_sequence_batch = jax.vmap(functools.partial(train.onehot, vocab_size=4))(
        np.array(
            [[0, 1, 0],
             [1, 0, 2],
             [1, 2, 0],
             [1, 2, 3]]
        )
    )
    np.testing.assert_equal(
        train.get_sequence_lengths(oh_sequence_batch, eos_id=0),
        np.array([1, 2, 3, 3], np.int32)
    )
    np.testing.assert_equal(
        train.get_sequence_lengths(oh_sequence_batch, eos_id=1),
        np.array([2, 1, 1, 1], np.int32)
    )
    np.testing.assert_equal(
        train.get_sequence_lengths(oh_sequence_batch, eos_id=2),
        np.array([3, 3, 2, 2], np.int32)
    )

  def test_mask_sequences(self):
    np.testing.assert_equal(
        train.mask_sequences(
            np.arange(1, 13).reshape((4, 3)),
            np.array([3, 2, 1, 0])
        ),
        np.array(
            [[1, 2, 3],
             [4, 5, 0],
             [7, 0, 0],
             [0, 0, 0]]
        )
    )

  def test_train_one_step(self):
    batch = train.get_batch(128)
    rng = random.PRNGKey(0)

    with nn.stochastic(rng):
      model = train.create_model()
      optimizer = train.create_optimizer(model, 0.003)
      optimizer, train_metrics = train.train_step(
          optimizer, batch, nn.make_rng())

    self.assertLessEqual(train_metrics['loss'], 5)
    self.assertGreaterEqual(train_metrics['accuracy'], 0)

  def test_decode_batch(self):
    with nn.stochastic(random.PRNGKey(0)):
      model = train.create_model()
      train.decode_batch(model, 5)

if __name__ == '__main__':
  absltest.main()