IPUStrategy

class pytorch_lightning.strategies.IPUStrategy(accelerator=None, device_iterations=1, autoreport=False, autoreport_dir=None, parallel_devices=None, cluster_environment=None, checkpoint_io=None, precision_plugin=None, training_opts=None, inference_opts=None)

Bases: pytorch_lightning.strategies.parallel.ParallelStrategy

Plugin for training on IPU devices.

Parameters

• device_iterations – Number of iterations to run on device at once before returning to host. This can be used as an optimization to speed up training. https://docs.graphcore.ai/projects/poptorch-user-guide/en/0.1.67/batching.html

• autoreport – Enable auto-reporting for IPUs using PopVision https://docs.graphcore.ai/projects/graphcore-popvision-user-guide/en/latest/graph/graph.html

• autoreport_dir – Optional directory to store autoReport output.

• training_opts – Optional poptorch.Options to override the default created options for training.

• inference_opts – Optional poptorch.Options to override the default created options for validation/testing and predicting.

all_gather(tensor, group=None, sync_grads=False)

Perform a all_gather on all processes.

Return type

Tensor

barrier(name=None)

Synchronizes all processes which blocks processes until the whole group enters this function.

Parameters

name (Optional[str]) – an optional name to pass into barrier.

Return type

None

batch_to_device(batch, device=None, dataloader_idx=0)

Moves the batch to the correct device.

The returned batch is of the same type as the input batch, just having all tensors on the correct device.

Parameters

• batch (Any) – The batch of samples to move to the correct device

• device (Optional[device]) – The target device

• dataloader_idx (int) – The index of the dataloader to which the batch belongs.

Return type

Any

broadcast(obj, src=0)

Broadcasts an object to all processes.

Parameters

• obj (object) – the object to broadcast

• src (int) – source rank

Return type

object

model_to_device()

Moves the model to the correct device.

Return type

None

on_predict_end()

Called when predict ends.

on_predict_start()

Called when predict begins.

on_test_end()

Called when test end.

on_test_start()

Called when test begins.

on_train_batch_start(batch, batch_idx)

Called in the training loop before anything happens for that batch.

Return type

None

on_train_end()

Called when train ends.

on_train_start()

Called when train begins.

on_validation_end()

Called when validation ends.

on_validation_start()

Called when validation begins.

predict_step(*args, **kwargs)

The actual predict step.

See predict_step() for more details

Return type

Union[Tensor, Dict[str, Any]]

reduce(tensor, *args, **kwargs)

Reduces the given tensor (e.g. across GPUs/processes).

Parameters

• tensor (Union[Tensor, Any]) – the tensor to sync and reduce

• group – the process group to reduce

• reduce_op – the reduction operation. Defaults to ‘mean’. Can also be a string ‘sum’ or ReduceOp.

Return type

Union[Tensor, Any]

setup(trainer)

Setup plugins for the trainer fit and creates optimizers.

Parameters

trainer (Trainer) – the trainer instance

Return type

None

setup_optimizers(trainer)

Creates optimizers and schedulers.

Parameters

trainer (Trainer) – the Trainer, these optimizers should be connected to

Return type

None

teardown()

This method is called to teardown the training process.

It is the right place to release memory and free other resources.

Return type

None

test_step(*args, **kwargs)

The actual test step.

See test_step() for more details

Return type

Union[Tensor, Dict[str, Any], None]

training_step(*args, **kwargs)

The actual training step.

See training_step() for more details

Return type

Union[Tensor, Dict[str, Any]]

validation_step(*args, **kwargs)

The actual validation step.

See validation_step() for more details

Return type

Union[Tensor, Dict[str, Any], None]

property is_global_zero: bool

Whether the current process is the rank zero process not only on the local node, but for all nodes.

Return type

bool

property root_device: torch.device

Return the root device.

Return type

device