pytorch_pfn_extras.runtime.PyTorchRuntime

class pytorch_pfn_extras.runtime.PyTorchRuntime(device_spec, options)

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

    If True, torch.cuda.amp.autocast is enabled. Default is False.

  • 'grad_scaler' (torch.cuda.amp.GradScaler):

    A gradient scaler that outputs are applied to.

Return type

None

__init__(device_spec, options)

Parameters

• device_spec (Union[str, torch.device]) –

• options (Dict[str, Any]) –

Return type

None

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.
move_module(module)	Transfers the module to the specific device.
move_tensor(tensor)	Transfers the tensor to the specific device.
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.

pytorch_pfn_extras.runtime.BaseRuntime pytorch_pfn_extras.training.ExtensionsManager