docstr: Trains a large network in a memory-efficient way.
docstr_tok: ["Trains","a","large","network","in","a","memory-efficient","way","."]
code: def test_train_memory_efficient(self):
      """Trains a large network in a memory-efficient way."""
      # This test requires > 16GB RAM, only run on TPUs. It does pass on GPU
      # and CPU when you run it locally, but it's too big for unit-testing.
      ram_limited = True  # Set to False to run this test locally.
      if fastmath.global_device_count() == 1 and ram_limited:
          return
  
      # Create the model.
      n_layers = (
          16  # 16 layers each 16K x 16K = 256M weights ~= 1GB, 16GB ram
      )
      model = tl.Serial(
          tl.Embedding(9, 16 * 1024),
          tl.Dup(),
          [
              [
                  tl.ReversibleHalfResidual(tl.Dense(16 * 1024)),
                  tl.ReversibleSwap(),
              ]
              for _ in range(n_layers)
          ],
          tl.Concatenate(),
          tl.Dense(9),
      )
  
      # Create inputs.
      inputs_batch = np.arange(8).reshape((2, 4))
      targets_batch = inputs_batch
      labeled_batch = (
          inputs_batch,
          targets_batch,
          np.ones_like(targets_batch),
      )
  
      def _data_gen():
          while True:
              yield labeled_batch
  
      # Run training.
      loss_layer = tl.WeightedCategoryCrossEntropy()
      task = training.TrainTask(
          _data_gen(), loss_layer, optimizers.Adafactor
      )
      eval_task = training.EvalTask(
          _data_gen(), [tl.WeightedCategoryCrossEntropy()]
      )
      loop = training.Loop(
          model,
          [task],
          eval_tasks=[eval_task],
          eval_at=lambda step_n: step_n == 2,
          use_memory_efficient_trainer=True,
      )
      self.assertEqual(0, loop.step)
      loop.run(n_steps=2)
      self.assertEqual(2, loop.step)
  
imports: import collections
  import os
  import time
  from absl.testing import absltest
  from jax import config
  import numpy
  from trax import data
  from trax import fastmath
  from trax import layers
  from trax import optimizers
  from trax import shapes
  from trax import test_utils
  from trax.layers import base
  from trax.models import transformer
  from trax.supervised import callbacks
  from trax.supervised import training
label_desc: Use image node_id 24 for calling the TrainingTest obj's underlying member method code with example usage: obj.test_train_memory_efficient() without return types