pytorch_pfn_extras.runtime.BaseRuntime#
- class pytorch_pfn_extras.runtime.BaseRuntime(device_spec, options)#
Bases:
object
A base class for collections of device-specific callback functions.
The function attributes of this class will be called from
ppe.to
orppe.handler.Handler
.ppe.runtime.runtime_registry
stores the runtime classes and dispatches them by feeding the corresponding name string as an input.- Parameters:
device_spec (torch.device or str) – The device that modules and tensors are transferred to.
options (dict) – A configuration dictionary that can be used from runtime method.
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
- convert_batch(args)#
Transfers the given batch to the specific device.
- Parameters:
args (object) – A batch data of any type.
- Returns:
A batch data transferred to the specific device of the same type as input.
- Return type:
Any
- 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