pytorch_pfn_extras.runtime.PyTorchRuntime#

class pytorch_pfn_extras.runtime.PyTorchRuntime(device_spec, options)#

Bases: BaseRuntime

A collections of callback functions for the devices that PyTorch supports by default.

Parameters:

device_spec (torch.device or str) – The device.
options (dict, optional) –
The configuration options.
- 'autocast' (bool or dict):
  If True, torch.cuda.amp.autocast is enabled. using {"enabled": True, "device_type": "cuda"} as autocast options. Default is False which corresponds to the following options {"enabled": False, "device_type": "cuda"} dict type. If dict, Options to pass to torch.autocast. Includes device_type, dtype among others.
- 'grad_scaler' (torch.cuda.amp.GradScaler):
  A gradient scaler that outputs are applied to.

Methods

`__init__`(device_spec, options)
`convert_batch`(args)	Transfers the given batch to the specific device.
`eval_post_step`(evaluator, module, batch_idx, ...)	The method called at the end of each evaluation.
`eval_pre_step`(evaluator, module, batch_idx, ...)	The method called at the beginning of each evaluation.
`execute`(code_block, batch)	Method called by the CodeBlocks API to do device dependent execution.
`initialize_module`(module, loader_or_batch[, ...])	Initializes the module at the beginning of training or inference.
`map`(func, iterable[, out_keys, device])	Method called by the user to apply function to iterable efficiently.
`move_module`(module)	Transfers the module to the specific device.
`move_tensor`(tensor)	Transfers the tensor to the specific device.
`trace`(event_name, arg)	Context manager for tracing PPE events in the custom device tools.
`train_cleanup`(module)	A method called only once when compleing a training run.
`train_epoch_begin`(module)	Preprocess of each epoch.
`train_epoch_end`(module)	Completion of each epoch.
`train_post_step`(trainer, module, batch_idx, ...)	Postprocess of each step.
`train_pre_step`(trainer, module, batch_idx, batch)	Preprocess of each step.
`train_validation_begin`(module)	The method called before each evaluation.
`train_validation_end`(module)	The method called after each evaluation.

__init__(device_spec, options)#

Parameters:

device_spec (Union[str, device]) –
options (Dict[str, Any]) –

Return type:

None

eval_post_step(evaluator, module, batch_idx, batch, outs)#

The method called at the end of each evaluation.

Parameters:

evaluator (Evaluator) – An evaluator.
module (torch.nn.Module) – A module.
batch_idx (int) – The batch index.
batch (list of torch.Tensor) – The list of input tensors of this batch.
outs (Any) – (list of torch.Tensor): The list of output tensors of this batch.

Return type:

None

Returns: None

eval_pre_step(evaluator, module, batch_idx, batch)#

The method called at the beginning of each evaluation.

Parameters:

evaluator (Evaluator) – An evaluator.
module (torch.nn.Module) – A module.
batch_idx (int) – The batch index.
batch (list of torch.Tensor) – The list of input tensors of this batch.

Return type:

None

Returns: None

execute(code_block, batch)#

Method called by the CodeBlocks API to do device dependent execution.

Parameters:

code_block (CodeBlock) – The codeblock requesting execution.
batch (dict of str, torch.Tensor) – The input tensors of this batch.

Returns:

The results of executing the codeblock on this runtime.

Return type:

Any

initialize_module(module, loader_or_batch, optimizer=None)#

Initializes the module at the beginning of training or inference.

Parameters:

module (torch.nn.Module) – A module.
loader_or_batch (DataLoader or torch.Tensor) – A data loader or a tensor.
optimizer (Optimizer or None) – An optimizer. This argument is sometimes used to copy LR from the original optimizer to the training model.

Return type:

None

Returns: None

map(func, iterable, out_keys=None, device='cpu')#

Method called by the user to apply function to iterable efficiently.

Parameters:

func (CodeBlock) – The function to be executed
iterable (Iterable[Any]) – The data
out_keys (Optional[Set[str]]) – The output keys that to be moved to the host device
device (Any) – The torch device that contains the final outputs

Returns:

The result of func

Return type:

Iterable[Any]

move_module(module)#

Transfers the module to the specific device.

Before this method is called, ppe.to will add this class as an new attribute (“_ppe_runtime”) to the input module.

Parameters:: module (torch.nn.Module) – A module.
Returns:: A module transferred to the specific device.
Return type:: Module

move_tensor(tensor)#

Transfers the tensor to the specific device.

Parameters:: tensor (torch.Tensor) – A tensor.
Returns:: A tensor transferred to the specific device.
Return type:: Tensor

classmethod trace(event_name, arg)#

Context manager for tracing PPE events in the custom device tools.

Parameters:

event_name (Optional[str]) – The name of the event being traced
arg (Any) – Custom argument for the tracer

Return type:

Generator[None, None, None]

train_cleanup(module)#

A method called only once when compleing a training run.

Parameters:: module (torch.nn.Module) – A module.
Return type:: None

Returns: None

train_epoch_begin(module)#

Preprocess of each epoch.

Parameters:: module (torch.nn.Module) – A module.
Return type:: None

Returns: None

train_epoch_end(module)#

Completion of each epoch.

Parameters:: module (torch.nn.Module) – A module.
Return type:: None

Returns: None

train_post_step(trainer, module, batch_idx, batch, outs)#

Postprocess of each step.

This method is called at the end of every steps: the set of (typically one) iterations and an update.

Parameters:

trainer (Trainer) – A trainer.
module (torch.nn.Module) – A module.
batch_idx (int) – The batch index.
batch (list of torch.Tensor) – The list of input tensors of this batch.
outs (Any) – (list of torch.Tensor): The list of output tensors of this batch.

Return type:

None

Returns: None

train_pre_step(trainer, module, batch_idx, batch)#

Preprocess of each step.

This method is called at the beginning of every steps: the set of (typically one) iterations and an update.

Parameters:

trainer (Trainer) – A trainer.
module (torch.nn.Module) – A module.
batch_idx (int) – The batch index.
batch (list of torch.Tensor) – The list of input tensors of this batch.

Return type:

None

Returns: None

train_validation_begin(module)#

The method called before each evaluation.

Parameters:: module (torch.nn.Module) – A module.
Return type:: None

Returns: None

train_validation_end(module)#

The method called after each evaluation.

Parameters:: module (torch.nn.Module) – A module.
Return type:: None

Returns: None