(output_dir,
          model=gin.REQUIRED,
          loss_fn=tl.WeightedCategoryCrossEntropy(),
          inputs=trax_inputs.batcher,
          optimizer=trax_opt.Adafactor,
          lr_schedule_fn=lr.multifactor,
          trainer_class=Trainer,
          steps=1000,
          checkpoints_at=None,
          permanent_checkpoints_at=None,
          eval_steps=10,
          eval_frequency=100,
          permanent_checkpoint_frequency=None,
          random_seed=None,
          save_graphs=True,
          metrics=None,
          checkpoint_highest=None,
          checkpoint_lowest=None,
          use_loop=True,
          loss_chunk_size=0,
          use_memory_efficient_trainer=False,
          adasum=False,
          init_checkpoint=None,
          callbacks=None,
          n_weights_shards=1,
          additional_train_tasks=None,
          additional_eval_tasks=None,
          additional_eval_streams=None)
return trainer.state
Train the model on the inputs.

  Args:
    output_dir: Directory where to put the logs and checkpoints.
    model: The model to train as a callable returning 2 callables, an init_fn
      and apply_fn.
    loss_fn: callable with signature: weights, trax.inputs.Inputs, model, state,
      rng -> loss.
    inputs: callable returning trax.inputs.Inputs.
    optimizer: The optimizer (see optimizers/base.py for signature).
    lr_schedule_fn: A learning rate schedule function, that when called returns
      a function from step to learning rate (a float).
    trainer_class: The trainer class to use.
    steps: int, total number of training steps.
    checkpoints_at: list of integers. Save a checkpoint for each training step
      in the list.
    permanent_checkpoints_at: list of integers. Save a permanent checkpoint for
      each training step in the list.
    eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
    eval_frequency: int, how often to run evaluation (every eval_frequency
      steps). If None or 0, eval disabled.
    permanent_checkpoint_frequency: int, how often to save permanent checkpoints
      (every permanent_checkpoint_frequency steps).
    random_seed: the random seed to use; time/os dependent if None (default).
    save_graphs: bool, if True, save computation graph to file.
    metrics: optionally override the default metrics dictionary.
    checkpoint_highest: save the checkpoint highest at this metric.
    checkpoint_lowest: save the checkpoint lowest at this metric.
    use_loop: whether to use training.Loop instead of Trainer.
    loss_chunk_size: int, if > 0 chunk loss into these sizes to save memory.
    use_memory_efficient_trainer: whether to use memory-efficient trainer.
    adasum: if True, use adaptive summation for multi-device gradients.
    init_checkpoint: a checkpoint for fine tuning.
    callbacks: a list of callbacks to call during training.
    n_weights_shards: shard weights into this many devices.
    additional_train_tasks: additional tasks which should be performed during
      training.
    additional_eval_tasks: additional tasks which should be performed during
      evaluation.
    additional_eval_streams: List[NamedStream], additional data streams that
      should be used during evaluation. Can be provided independently of
      additional_eval_tasks.

  Returns:
    trax.TrainerState or training.Loop if use_loop is True
Train the model on the inputs.
Train the model on the inputs .
def train(output_dir,
          model=gin.REQUIRED,
          loss_fn=tl.WeightedCategoryCrossEntropy(),
          inputs=trax_inputs.batcher,
          optimizer=trax_opt.Adafactor,
          lr_schedule_fn=lr.multifactor,
          trainer_class=Trainer,
          steps=1000,
          checkpoints_at=None,
          permanent_checkpoints_at=None,
          eval_steps=10,
          eval_frequency=100,
          permanent_checkpoint_frequency=None,
          random_seed=None,
          save_graphs=True,
          metrics=None,
          checkpoint_highest=None,
          checkpoint_lowest=None,
          use_loop=True,
          loss_chunk_size=0,
          use_memory_efficient_trainer=False,
          adasum=False,
          init_checkpoint=None,
          callbacks=None,
          n_weights_shards=1,
          additional_train_tasks=None,
          additional_eval_tasks=None,
          additional_eval_streams=None):
  """Train the model on the inputs.

  Args:
    output_dir: Directory where to put the logs and checkpoints.
    model: The model to train as a callable returning 2 callables, an init_fn
      and apply_fn.
    loss_fn: callable with signature: weights, trax.inputs.Inputs, model, state,
      rng -> loss.
    inputs: callable returning trax.inputs.Inputs.
    optimizer: The optimizer (see optimizers/base.py for signature).
    lr_schedule_fn: A learning rate schedule function, that when called returns
      a function from step to learning rate (a float).
    trainer_class: The trainer class to use.
    steps: int, total number of training steps.
    checkpoints_at: list of integers. Save a checkpoint for each training step
      in the list.
    permanent_checkpoints_at: list of integers. Save a permanent checkpoint for
      each training step in the list.
    eval_steps: int, num of steps per evaluation. If None or 0, eval disabled.
    eval_frequency: int, how often to run evaluation (every eval_frequency
      steps). If None or 0, eval disabled.
    permanent_checkpoint_frequency: int, how often to save permanent checkpoints
      (every permanent_checkpoint_frequency steps).
    random_seed: the random seed to use; time/os dependent if None (default).
    save_graphs: bool, if True, save computation graph to file.
    metrics: optionally override the default metrics dictionary.
    checkpoint_highest: save the checkpoint highest at this metric.
    checkpoint_lowest: save the checkpoint lowest at this metric.
    use_loop: whether to use training.Loop instead of Trainer.
    loss_chunk_size: int, if > 0 chunk loss into these sizes to save memory.
    use_memory_efficient_trainer: whether to use memory-efficient trainer.
    adasum: if True, use adaptive summation for multi-device gradients.
    init_checkpoint: a checkpoint for fine tuning.
    callbacks: a list of callbacks to call during training.
    n_weights_shards: shard weights into this many devices.
    additional_train_tasks: additional tasks which should be performed during
      training.
    additional_eval_tasks: additional tasks which should be performed during
      evaluation.
    additional_eval_streams: List[NamedStream], additional data streams that
      should be used during evaluation. Can be provided independently of
      additional_eval_tasks.

  Returns:
    trax.TrainerState or training.Loop if use_loop is True
  """
  base.N_WEIGHTS_SHARDS = n_weights_shards
  if (permanent_checkpoint_frequency is not None
      and permanent_checkpoints_at is not None):
    raise ValueError('Only one of ["permanent_checkpoint_frequency", '
                     '"permanent_checkpoints_at"] should be set.')
  if use_loop:
    n_devices = num_devices() or fastmath.local_device_count()

    # Prepare the training task.
    # Inputs is either an Inputs instance or a function that returns it.
    if callable(inputs):  # If we pass a function, e.g., through gin, call it.
      inputs = inputs()
    opt = optimizer if use_memory_efficient_trainer else optimizer()
    train_task = training.TrainTask(
        inputs.train_stream(n_devices),
        loss_layer=loss_fn,
        optimizer=opt,
        lr_schedule=lr_schedule_fn(),
        n_steps_per_checkpoint=eval_frequency,
        n_steps_per_permanent_checkpoint=permanent_checkpoint_frequency)

    if additional_train_tasks is None:
      additional_train_tasks = []

    # Prepare the evaluation.
    metrics_dict = metrics if metrics is not None else _DEFAULT_METRICS
    names, metrics = zip(*metrics_dict.items())
    eval_task = training.EvalTask(inputs.eval_stream(n_devices),
                                  metrics,
                                  metric_names=names,
                                  n_eval_batches=eval_steps)

    if additional_eval_tasks is None:
      additional_eval_tasks = []

    additional_eval_tasks_from_streams = []
    if additional_eval_streams is not None:
      for stream in additional_eval_streams:
        additional_eval_tasks_from_streams.append(
            training.EvalTask(stream.stream,
                              metrics,
                              metric_names=names,
                              n_eval_batches=eval_steps,
                              export_prefix=stream.name))

    # Prepare the training loop.
    checkpoint_at = None
    if checkpoints_at is not None:
      checkpoint_at = lambda step: step in checkpoints_at
    permanent_checkpoint_at = None
    if permanent_checkpoints_at is not None:
      permanent_checkpoint_at = (lambda step: step in permanent_checkpoints_at)

    # Setup the model.
    model_train = model(mode='train')
    model_predict_eval = model(mode='eval')
    if init_checkpoint:
      model_train.init_from_file(init_checkpoint, weights_only=True)
      model_predict_eval.init_from_file(init_checkpoint, weights_only=True)
    loop = training.Loop(
        model_train, [train_task] + additional_train_tasks,
        eval_model=model_predict_eval,
        eval_tasks=[eval_task] +
        additional_eval_tasks + additional_eval_tasks_from_streams,
        output_dir=output_dir,
        checkpoint_at=checkpoint_at,
        checkpoint_low_metric=checkpoint_lowest,
        checkpoint_high_metric=checkpoint_highest,
        permanent_checkpoint_at=permanent_checkpoint_at,
        n_devices=n_devices,
        loss_chunk_size=loss_chunk_size,
        use_memory_efficient_trainer=use_memory_efficient_trainer,
        adasum=adasum,
        random_seed=random_seed,
        callbacks=callbacks,
    )

    steps_to_go = steps - loop.step
    if steps_to_go <= 0:
      log('Stop training, already reached the total training steps %d' % steps)
      return loop

    # Train and return the loop.
    loop.run(steps_to_go)
    return loop

  n_devices = num_devices()
  trainer = trainer_class(model, loss_fn, optimizer, lr_schedule_fn(), inputs,
                          output_dir,
                          random_seed=random_seed,
                          n_devices=n_devices,
                          checkpoints_at=checkpoints_at,
                          metrics=metrics,
                          checkpoint_lowest=checkpoint_lowest,
                          checkpoint_highest=checkpoint_highest,
                          init_checkpoint=init_checkpoint)

  epoch_steps = [steps]  # Only training if eval_frequency is 0 or None
  if eval_frequency and eval_steps > 0:
    epoch_steps = itertools.chain([1,  # first epoch only 1 step
                                   eval_frequency - 1],
                                  itertools.repeat(eval_frequency))
  trainer.log_step('Starting training using %d devices' % trainer.n_devices)
  trainer.print_n_weights()

  try:
    for epoch_steps in epochs(steps, trainer.step, epoch_steps):
      trainer.train_epoch(epoch_steps, eval_steps)

      # Bookkeeping we do at the first step
      if trainer.step == 1:
        # Save computation graph (single-device only for now)
        if (save_graphs and fastmath.is_backend(fastmath.Backend.JAX)):
          trainer.save_computation_graphs()

        # Save Gin config
        trainer.save_gin()

    trainer.log_step('Training done')
  except Exception as e:
    raise e
  finally:
    trainer.close()
  return trainer.state
def train ( output_dir , model = gin . REQUIRED , loss_fn = tl . WeightedCategoryCrossEntropy ( ) , inputs = trax_inputs . batcher , optimizer = trax_opt . Adafactor , lr_schedule_fn = lr . multifactor , trainer_class = Trainer , steps = 1000 , checkpoints_at = None , permanent_checkpoints_at = None , eval_steps = 10 , eval_frequency = 100 , permanent_checkpoint_frequency = None , random_seed = None , save_graphs = True , metrics = None , checkpoint_highest = None , checkpoint_lowest = None , use_loop = True , loss_chunk_size = 0 , use_memory_efficient_trainer = False , adasum = False , init_checkpoint = None , callbacks = None , n_weights_shards = 1 , additional_train_tasks = None , additional_eval_tasks = None , additional_eval_streams = None ) : base . N_WEIGHTS_SHARDS = n_weights_shards if ( permanent_checkpoint_frequency is not None and permanent_checkpoints_at is not None ) : raise ValueError ( 'Only one of ["permanent_checkpoint_frequency", ' '"permanent_checkpoints_at"] should be set.' ) if use_loop : n_devices = num_devices ( ) or fastmath . local_device_count ( ) # Prepare the training task. # Inputs is either an Inputs instance or a function that returns it. if callable ( inputs ) : # If we pass a function, e.g., through gin, call it. inputs = inputs ( ) opt = optimizer if use_memory_efficient_trainer else optimizer ( ) train_task = training . TrainTask ( inputs . train_stream ( n_devices ) , loss_layer = loss_fn , optimizer = opt , lr_schedule = lr_schedule_fn ( ) , n_steps_per_checkpoint = eval_frequency , n_steps_per_permanent_checkpoint = permanent_checkpoint_frequency ) if additional_train_tasks is None : additional_train_tasks = [ ] # Prepare the evaluation. metrics_dict = metrics if metrics is not None else _DEFAULT_METRICS names , metrics = zip ( * metrics_dict . items ( ) ) eval_task = training . EvalTask ( inputs . eval_stream ( n_devices ) , metrics , metric_names = names , n_eval_batches = eval_steps ) if additional_eval_tasks is None : additional_eval_tasks = [ ] additional_eval_tasks_from_streams = [ ] if additional_eval_streams is not None : for stream in additional_eval_streams : additional_eval_tasks_from_streams . append ( training . EvalTask ( stream . stream , metrics , metric_names = names , n_eval_batches = eval_steps , export_prefix = stream . name ) ) # Prepare the training loop. checkpoint_at = None if checkpoints_at is not None : checkpoint_at = lambda step : step in checkpoints_at permanent_checkpoint_at = None if permanent_checkpoints_at is not None : permanent_checkpoint_at = ( lambda step : step in permanent_checkpoints_at ) # Setup the model. model_train = model ( mode = 'train' ) model_predict_eval = model ( mode = 'eval' ) if init_checkpoint : model_train . init_from_file ( init_checkpoint , weights_only = True ) model_predict_eval . init_from_file ( init_checkpoint , weights_only = True ) loop = training . Loop ( model_train , [ train_task ] + additional_train_tasks , eval_model = model_predict_eval , eval_tasks = [ eval_task ] + additional_eval_tasks + additional_eval_tasks_from_streams , output_dir = output_dir , checkpoint_at = checkpoint_at , checkpoint_low_metric = checkpoint_lowest , checkpoint_high_metric = checkpoint_highest , permanent_checkpoint_at = permanent_checkpoint_at , n_devices = n_devices , loss_chunk_size = loss_chunk_size , use_memory_efficient_trainer = use_memory_efficient_trainer , adasum = adasum , random_seed = random_seed , callbacks = callbacks , ) steps_to_go = steps - loop . step if steps_to_go <= 0 : log ( 'Stop training, already reached the total training steps %d' % steps ) return loop # Train and return the loop. loop . run ( steps_to_go ) return loop n_devices = num_devices ( ) trainer = trainer_class ( model , loss_fn , optimizer , lr_schedule_fn ( ) , inputs , output_dir , random_seed = random_seed , n_devices = n_devices , checkpoints_at = checkpoints_at , metrics = metrics , checkpoint_lowest = checkpoint_lowest , checkpoint_highest = checkpoint_highest , init_checkpoint = init_checkpoint ) epoch_steps = [ steps ] # Only training if eval_frequency is 0 or None if eval_frequency and eval_steps > 0 : epoch_steps = itertools . chain ( [ 1 , # first epoch only 1 step eval_frequency - 1 ] , itertools . repeat ( eval_frequency ) ) trainer . log_step ( 'Starting training using %d devices' % trainer . n_devices ) trainer . print_n_weights ( ) try : for epoch_steps in epochs ( steps , trainer . step , epoch_steps ) : trainer . train_epoch ( epoch_steps , eval_steps ) # Bookkeeping we do at the first step if trainer . step == 1 : # Save computation graph (single-device only for now) if ( save_graphs and fastmath . is_backend ( fastmath . Backend . JAX ) ) : trainer . save_computation_graphs ( ) # Save Gin config trainer . save_gin ( ) trainer . log_step ( 'Training done' ) except Exception as e : raise e finally : trainer . close ( ) return trainer . state