def test_train_save_restore_transformer(self):
    """Saves and restores a checkpoint to check for equivalence."""
    vocab_size = 8
    task = training.TrainTask(
        _very_simple_transformer_data(), tl.L2Loss(), optimizers.SGD(.01))
    eval_task = training.EvalTask(
        _very_simple_transformer_data(),  # deliberately re-using training data
        [tl.L2Loss()],
        metric_names=['SGD.L2Loss'])
    tmp_dir = self.create_tempdir().full_path

    def _make_model_and_session():
      m = transformer.TransformerLM(
          vocab_size, d_model=4, d_ff=4, n_layers=1, n_heads=2, dropout=0.)
      ts = training.Loop(m, [task], eval_tasks=[eval_task],
                         eval_at=lambda step_n: step_n % 2 == 0,
                         output_dir=tmp_dir)
      return m, ts

    model, training_session = _make_model_and_session()
    self.assertEqual(0, training_session.step)
    training_session.run(n_steps=1)
    training_session.save_checkpoint('model')
    model2, training_session2 = _make_model_and_session()

    x = np.ones((2, 2)).astype(np.int32)
    y1 = model(x, rng=fastmath.random.get_prng(0))
    y2 = model2(x, rng=fastmath.random.get_prng(0))
    self.assertEqual(str(y1), str(y2))

    training_session2.run(n_steps=1)
    y1 = model(x, rng=fastmath.random.get_prng(0))
    y2 = model2(x, rng=fastmath.random.get_prng(0))
    self.assertNotEqual(str(y1), str(y2))
Saves and restores a checkpoint to check for equivalence.
trax/supervised/training_test.py