composer.core.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.

property is_model_deepspeed[source]#

Whether model is an instance of a DeepSpeedEngine.

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.)