state#

The state of the trainer.

Classes

State

The state of the trainer.

class composer.core.state.State(model, rank_zero_seed, run_name, max_duration=None, grad_accum=1, train_dataloader=None, evaluators=None, dataloader=None, dataloader_label=None, dataloader_len=- 1, precision=Precision.FP32, optimizers=None, scaler=None, algorithms=None, callbacks=None, deepspeed_config=None)[source]#

Bases: composer.core.serializable.Serializable

The state of the trainer.

Contains variables that the trainer tracks throughout the training loop. Note that all the necessary parts (i.e., serialized_attributes) of state are serialized when the trainer is checkpointed so that it can be used restore the trainer and continue training from a checkpoint. algorithms are able to modify an instance of this class in-place.

Note

An instance of this class is automatically constructed by the Trainer constructor. A user need not instantiate this class.

Parameters

model (Module) – The model, typically as a subclass of ComposerModel.
rank_zero_seed (int) – The seed used on the rank zero process. It is assumed that each rank’s seed is rank_zero_seed + dist.get_global_rank().
run_name (str) – The name for this training run.
grad_accum (int, optional) – The number of gradient accumulation steps to use. With this argument, micro batch size for each device becomes microbatch_size = train_batch_size / (num_devices * grad_accum).
train_dataloader (types.DataLoader, optional) – Dataloader used for training
evaluators (Evalutor | Evaluators, optional) – Evaluator used for evaluation.
dataloader (types.DataLoader, optional) – The active DataLoader.
dataloader_len (int | Time[int], optional) – The number of batches per dataloader iteration (e.g. epoch). The trainer will yield the first dataloader_len batches per iteration. If -1 (the default), the entire dataloader will be iterated over.
dataloader_label (str, optional) –
The name for the dataloader. Required if dataloader is specified. (default: None)

By convention, the training dataloader is called 'train'. The evaluator dataloader is called 'eval', or when multiple evaluators are used, the name of the evaluator.
max_duration (str | Time, optional) – The maximum duration to train for. (default: None)
precision (str | Precision) – The numerical precision to use for training. See Precision for the supported precisions.
optimizers (Optimizer | Sequence[Optimizer], optional) – The optimizer being used to train the model. Multiple optimizers are not currently supported.
schedulers (types.PyTorchScheduler | Sequence[types.PyTorchScheduler], optional) – The learning rate scheduler (can also be a list or tuple of schedulers).
scaler (GradScaler, optional) – The gradient scaler in use for mixed precision training.
algorithms (Algorithm | Sequence[Algorithm], optional) – The algorithms used for training.
callbacks (Callback | Sequence[Callback], optional) – The callbacks used for training.
deepspeed_config (Dict[str, Any], optional) – The configuration dictionary for deepspeed.

batch#

The batch. This will be the entire batch during the Event.AFTER_DATALOADER, or a microbatch between Event.BATCH_START and Event.BATCH_END.

Type: types.Batch

current_metrics#

The current computed metrics, organized by dataloader label and then by metric name. The train dataloader is labeled 'train'. If not using an Evaluator, the eval dataloader is labeled 'eval'. Otherwise, the evaluator label is used.

For example:

>>> trainer = Trainer(
...     ...,
...     compute_training_metrics=True,
...     train_dataloader=train_dataloader,
...     eval_dataloader=eval_dataloader,
... )
>>> trainer.fit()
>>> trainer.state.current_metrics
{'train': {'Accuracy': tensor(...)}, 'eval': {'CrossEntropy': tensor(...), 'Accuracy': tensor(...)}}

Or, when using an Evaluator:

>>> from torchmetrics import Accuracy
>>> from composer.core import Evaluator
>>> trainer = Trainer(
...     ...,
...     compute_training_metrics=True,
...     train_dataloader=train_dataloader,
...     eval_dataloader=[
...         Evaluator(label='eval1', dataloader=eval_1_dl, metrics=Accuracy()),
...         Evaluator(label='eval2', dataloader=eval_2_dl, metrics=Accuracy()),
...     ],
... )
>>> trainer.fit()
>>> trainer.state.current_metrics
{'train': {'Accuracy': tensor(...)}, 'eval1': {'Accuracy': tensor(...)}, 'eval2': {'Accuracy': tensor(...)}}

Type: Dict[str, Dict[str, Any]]

eval_timestamp#

The timestamp for the current evaluation dataloader. This timestamp is reset before the dataloader is evaluated. The epoch attribute for this timestamp is always 0.

Type: Timestamp

grad_accum#

The number of gradient accumulation steps per batch.

Type: int

loss#

The most recently computed loss.

Type: Tensor | Sequence[Tensor]

model#

The training model.

Note

When using DeepSpeed or multi-rank training, the model will be wrapped with DeepSpeedEngine or DistributedDataParallel, respectively.

Type: Module

outputs#

The most recently computed output from the model’s forward pass.

Type: Tensor | Sequence[Tensor]

predict_timestamp#

The timestamp for the current prediction dataloader. This timestamp is reset before the dataloader is used. The epoch attribute for this timestamp is always 0.

Type: Timestamp

profiler#

The profiler (if profiling is enabled), or None if not profiling.

Type: Profiler

rank_zero_seed#

The seed of the rank zero process.

Type: int

run_name#

The name for this training run.

Type: str

scaler#

The gradient scaler if using mixed-precision training, or None if not using mixed-precision training.

Type: GradScaler

serialized_attributes#

The names of the attribute which are serialized in a checkpoint.

By default, the following attributes are serialized:

Attribute	Description
model	The model under training.
optimizers	The optimizers being used to train the model.
schedulers	The learning rate schedulers.
algorithms	The algorithms used for training.
callbacks	The callbacks used for training.
scaler	The gradient scaler in use for mixed precision training.
timestamp	The timestamp that tracks training loop progress.
rank_zero_seed	The seed of the rank zero process.
current_metrics	The current metrics.
run_name	The run name for training.

Type: List[str]

timestamp#

The current training timestamp.

Type: Timestamp

train_dataloader#

The training dataloader. (May be None if not training.)

Type: Iterable

property algorithms[source]#: The algorithms.

batch_get_item(key)[source]#

Gets element from batch either specified by key or user-specified function.

See batch_get in utils/batch_helpers.py for examples.

Parameters: key (str | int | Tuple[Callable, Callable] | Any, optional) – A key to index into the batch or a user-specified function to do the extracting. A pair of callables is also supported for cases where a get and set function pair are both passed (like in Algorithms). The getter is assumed to be the first of the pair.
Returns: The part of the batch specified by the key. This could be any type – depending on what the batch is composed of.

batch_set_item(key, value)[source]#

Sets the element specified by the key of the set_fn to the specified value.

This is not an in-place operation, as for tuple-typed batches, a new batch object must be created to modify them.

See batch_set in utils/batch_helpers.py for examples.

Parameters

key (str | int | Tuple[Callable, Callable] | Any, optional) – A key to index into the batch or a user-specified function to do the setting. A pair of callables is also supported for cases where a get and set function pair are both passed (like in Algorithms). The setter is assumed to be the second of the pair.
value (Any) – The value that batch[key] or batch.key gets set to or that the user-defined set function sets a part of the batch to.

Returns

batch (Any) – The updated batch with value set at key.

property callbacks[source]#: The callbacks.

property dataloader[source]#: The active dataloader.

property dataloader_label[source]#

The dataloader label for the active dataloader.

By default, the training dataloader is called 'train'. The evaluator dataloader is called 'eval', or when multiple evaluators are used, the name of the evaluator. However, the dataloader label can be explicitely specified in Trainer.fit() and Trainer.eval().

Returns: Optional[str] – The dataloader label, or None if no dataloader is set.

property dataloader_len[source]#

The number of batches per dataloader iteration (e.g. epoch), as used by the trainer.

Note

If not explicitely specified, this value is an approximation, as it depends on len(self.dataloader). See the PyTorch DataLoader Documentation for more information.

Returns

Optional[Time[int]] – The number of batches per dataloader iteration (e.g. epoch), or None if no dataloader
is defined or if the dataloader has an unknown length (e.g. streaming dataloaders)

property deepspeed_enabled[source]#: Indicates if deepspeed is enabled.

property deepspeed_model[source]#: Cast model to DeepSpeedEngine.

property evaluators[source]#: The evaluators.

get_elapsed_duration()[source]#

Get the elapsed training duration.

Returns: Optional[Time[float]] – The elapsed duration, in TimeUnit.DURATION. Time(0.0, TimeUnit.DURATION) represents the beginning of training and Time(1.0, TimeUnit.DURATION) represents a completed training process. Returns None if max_duration is None.

property is_model_ddp[source]#: Whether model is an instance of a DistributedDataParallel.

load_model_state(state_dict, strict)[source]#

Loads the model’s state from a state_dict.

Parameters

state_dict (Dict[str, Any]) – The state dict, generated from a previous call to state_dict().
strict (bool) – Whether the keys (i.e., model parameter names) in the model state dict should perfectly match the keys in the model instance.

load_state_dict(state, strict=False)[source]#

Loads the state.

Parameters

state (Dict[str, Any]) – object returned from call to state_dict().
strict (bool) – whether the keys in the state["model"] should perfectly match the keys in the self.model. Defaults to False.

property max_duration[source]#: The maximum training duration.

property optimizers[source]#: The optimizers.

property precision[source]#

The numerical precision to use for training.

See Precision for the supported precisions.

property schedulers[source]#: The schedulers.

property seed[source]#: The seed for the current rank.

set_dataloader(dataloader=None, dataloader_label=None, dataloader_len=- 1)[source]#

Update the active dataloader and dataloader label.

Parameters

dataloader (Iterable, optional) – The dataloader. Defaults to None.
dataloader_label (str, optional) – The dataloader label. Must be None if and only if dataloader is None. Defaults to None.
dataloader_len (int, int) – The number of batches per dataloader iteration (e.g. epoch), as used by the trainer. Set to -1 to iterate over the entire dataset. (Default: -1.)