docstr: Train two-head MNIST model a bit, to compare to other implementations.
docstr_tok: ["Train","two-head","MNIST","model","a","bit",",","to","compare","to","other","implementations","."]
code: def test_train_mnist_multitask(self, mock_stdout):
      """Train two-head MNIST model a bit, to compare to other implementations."""
      mnist_model = _build_model(two_heads=True)
      # MNIST classification task.
      (cls_task, cls_eval_task) = _mnist_tasks(
          head=tl.Select([0], n_in=2)
      )
      # Auxiliary brightness prediction task.
      reg_task = training.TrainTask(
          itertools.cycle(_mnist_brightness_dataset().train_stream(1)),
          tl.Serial(tl.Select([1]), tl.L2Loss()),
          adam.Adam(0.001),
      )
      reg_eval_task = training.EvalTask(
          itertools.cycle(_mnist_brightness_dataset().eval_stream(1)),
          [tl.Serial(tl.Select([1]), tl.L2Loss())],
          n_eval_batches=1,
          metric_names=["L2"],
      )
      training_session = training.Loop(
          mnist_model,
          tasks=[cls_task, reg_task],
          eval_tasks=[cls_eval_task, reg_eval_task],
          eval_at=lambda step_n: step_n % 20 == 0,
          which_task=lambda step_n: step_n % 2,
      )
  
      training_session.run(n_steps=100)
      self.assertEqual(training_session.step, 100)
  
      # Assert that we reach at least 80% eval accuracy on MNIST.
      self.assertGreater(
          _read_metric("WeightedCategoryAccuracy", mock_stdout), 0.8
      )
      # Assert that we get below 0.03 brightness prediction error.
      self.assertLess(_read_metric("L2", mock_stdout), 0.03)
  
file_path: trax/trax/supervised/mnist_test.py
imports: import io
  import itertools
  from unittest import mock
  from absl.testing import absltest
  from trax import layers
  from trax.data import inputs
  from trax.data import tf_inputs
  from trax.optimizers import adam
  from trax.supervised import training
label_desc: Use image node_id 2 for calling the MnistTest obj's underlying member method code with example usage: obj.test_train_mnist_multitask(mock_stdout) without return types