orbit.StandardTrainer

Implements standard functionality on top of the AbstractTrainer API.

Inherits From: AbstractTrainer

orbit.StandardTrainer(
    train_dataset,
    options: Optional[orbit.StandardTrainerOptions] = None
)

This class structures the training "inner loop" roughly as follows:

train_loop_begin()
for _ in range(num_steps):
  train_step(train_iterator)
return train_loop_end()

Calls to train_loop_begin and train_loop_end are always done in eager mode, while the loop/train_step may be implemented using tf.while and/or tf.function, as determined by the options passed to __init__.

Args
train_dataset	A tf.nest-compatible structure of tf.data.Dataset or DistributedDataset.
options	An orbit.StandardTrainerOptions instance.

Attributes
name	Returns the name of this module as passed or determined in the ctor. Note: This is not the same as the self.name_scope.name which includes parent module names.
name_scope	Returns a tf.name_scope instance for this class.
non_trainable_variables	Sequence of non-trainable variables owned by this module and its submodules. Note: this method uses reflection to find variables on the current instance and submodules. For performance reasons you may wish to cache the result of calling this method if you don't expect the return value to change.
submodules	Sequence of all sub-modules. Submodules are modules which are properties of this module, or found as properties of modules which are properties of this module (and so on). a = tf.Module() b = tf.Module() c = tf.Module() a.b = b b.c = c list(a.submodules) == [b, c] True list(b.submodules) == [c] True list(c.submodules) == [] True
train_dataset	The current training dataset.
trainable_variables	Sequence of trainable variables owned by this module and its submodules. Note: this method uses reflection to find variables on the current instance and submodules. For performance reasons you may wish to cache the result of calling this method if you don't expect the return value to change.
variables	Sequence of variables owned by this module and its submodules. Note: this method uses reflection to find variables on the current instance and submodules. For performance reasons you may wish to cache the result of calling this method if you don't expect the return value to change.

Methods

create_train_loop_fn

View source

create_train_loop_fn()

Creates a training loop from the current step function and options.

Returns
The train loop function, i.e. wrapper of multiple train steps.

train

View source

train(
    num_steps: tf.Tensor
) -> Optional[runner.Output]

Implements num_steps steps of training.

Args
num_steps	The number of training steps to run. This corresponds directly to the number of calls made to train_step.

Returns
The output of train_loop_end.

train_loop_begin

View source

train_loop_begin()

Called once at the beginning of the training loop.

This method is always called in eager mode, and is a good place to reset metrics that accumulate values over multiple steps of training.

Note that this method is called before dataset iterator creation.

train_loop_end

View source

train_loop_end() -> Optional[runner.Output]

Called once at the end of the training loop.

This method is always called in eager mode, and is a good place to get metric results. The value returned from this function will be returned as-is from the train method implementation provided by StandardTrainer.

Returns
The function may return a dictionary of Tensors, which will be written to logs and as TensorBoard summaries. It can also be a nested dictionary, yielding a hierarchy of summary directories.

train_step

View source

@abc.abstractmethod
train_step(
    iterator
)

Implements one step of training.

What a "step" consists of is up to the implementer. When using distribution strategies, the call to this method takes place in the "cross-replica context" for generality, to allow e.g. multiple iterator dequeues and calls to strategy.run.

Note that if use_tf_function=True, all the code inside train_step should be compatible with tf.function tracing (and in particular, any state modifications involving self should be avoided). In some cases, non- tf.function compatible code can be moved to train_loop_begin or train_loop_end, which always execute eagerly.

Args
iterator	A tf.nest-compatible structure of tf.data.Iterator or DistributedIterator. The structure of this input matches the structure of train_dataset as passed to __init__.

with_name_scope

@classmethod
with_name_scope(
    method
)

Decorator to automatically enter the module name scope.

class MyModule(tf.Module):
  @tf.Module.with_name_scope
  def __call__(self, x):
    if not hasattr(self, 'w'):
      self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
    return tf.matmul(x, self.w)

Using the above module would produce tf.Variables and tf.Tensors whose names included the module name:

mod = MyModule()
mod(tf.ones([1, 2]))
<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
mod.w
<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
numpy=..., dtype=float32)>

Args
method	The method to wrap.

Returns
The original method wrapped such that it enters the module's name scope.