XLAStrategy

class lightning.pytorch.strategies.XLAStrategy(accelerator=None, parallel_devices=None, checkpoint_io=None, precision_plugin=None, debug=False, **_)

기반 클래스: lightning.pytorch.strategies.ddp.DDPStrategy

Strategy for training multiple TPU devices using the torch_xla.distributed.xla_multiprocessing.spawn() method.

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

Function to gather a tensor from several distributed processes.

매개변수

• tensor (Tensor) – tensor of shape (batch, …)

• group (Optional[Any]) – not available with TPUs

• sync_grads (bool) – flag that allows users to synchronize gradients for the all_gather operation

반환 형식

Tensor

반환

A tensor of shape (world_size, batch, …)

barrier(name=None, *args, **kwargs)

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

매개변수

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

반환 형식

None

broadcast(obj, src=0)

Broadcasts an object to all processes.

매개변수

• obj (TypeVar(TBroadcast)) – the object to broadcast

• src (int) – source rank

반환 형식

TypeVar(TBroadcast)

connect(model)

Called by the accelerator to connect the accelerator and the model with this plugin.

반환 형식

None

model_to_device()

Moves the model to the correct device.

반환 형식

None

on_train_batch_start(batch, batch_idx)

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

반환 형식

None

predict_step(*args, **kwargs)

The actual predict step.

See predict_step() for more details

반환 형식

Union[Tensor, Dict[str, Any]]

process_dataloader(dataloader)

Wraps the dataloader if necessary.

매개변수

dataloader (object) – iterable. Ideally of type: torch.utils.data.DataLoader

반환 형식

None

reduce(output, group=None, reduce_op=None)

Reduces a tensor from several distributed processes to one aggregated tensor.

매개변수

• tensor – the tensor to sync and reduce

• group (Optional[Any]) – the process group to gather results from. Defaults to all processes (world)

• reduce_op (Union[ReduceOp, str, None]) – the reduction operation. Defaults to ‘mean’/’avg’. Can also be a string ‘sum’ to calculate the sum during reduction.

반환 형식

Tensor

반환

reduced value, except when the input was not a tensor the output remains is unchanged

remove_checkpoint(filepath)

Remove checkpoint filepath from the filesystem.

매개변수

filepath (Union[str, Path]) – Path to checkpoint

반환 형식

None

save_checkpoint(checkpoint, filepath, storage_options=None)

Save model/training states as a checkpoint file through state-dump and file-write.

매개변수

• checkpoint (Dict[str, Any]) – dict containing model and trainer state

• filepath (Union[str, Path]) – write-target file’s path

• storage_options (Optional[Any]) – parameter for how to save to storage, passed to CheckpointIO plugin

반환 형식

None

setup(trainer)

Setup plugins for the trainer fit and creates optimizers.

매개변수

trainer (Trainer) – the trainer instance

반환 형식

None

teardown()

This method is called to teardown the training process.

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

반환 형식

None

test_step(*args, **kwargs)

The actual test step.

See test_step() for more details

반환 형식

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

validation_step(*args, **kwargs)

The actual validation step.

See validation_step() for more details

반환 형식

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

property root_device: torch.device

Return the root device.

반환 형식

device