Pipeline Parallelism#

vision_architectures.utils.pipeline_parallelism.get_device(device)[source]#

Convert to torch.device object.

Return type:

device

vision_architectures.utils.pipeline_parallelism.move_to_device(data, device)[source]#

Move data to the specified device.

Parameters:

  • data – The data to move.

  • device – The device to move the data to.

Returns:

The data moved to the specified device.

vision_architectures.utils.pipeline_parallelism.parallelize_module(module, processing_device, output_device=None)[source]#
vision_architectures.utils.pipeline_parallelism.unparallelize_module(module, common_device=None)[source]#

Unparallelize a module that was parallelized with parallelize_module. This will move the module to the common device and restore the original forward and extra_repr functions.

vision_architectures.utils.pipeline_parallelism.parallelize_pipeline(model, module_to_device)[source]#

Parallelize a model across multiple devices. It is recommended to pass a root device too i.e. {“”: DEVICE, …} to avoid errors in the calculations of the outermost modules.

Parameters:

  • model (Module) – The model to parallelize. An unparallelized model is expected, otherwise you will face recursion errors.

  • module_to_device (dict[str, device | str | list[device | str]]) – A dictionary mapping module names to devices. Keys are modules names with nested modules separated by dots (e.g., “module.submodule”). Note that the parallelism is performed using Level Order Traversal (i.e. BFS) of the model. Therefore the device of the deepest module in the dictionary will be overwritten even if it’s parent is also specified in the dictionary. Value is either a device or a 2-tuple of devices. The first device is the processing device, and the second device is the output device.

Return type:

Module

Returns:

The parallelized pipeline.

vision_architectures.utils.pipeline_parallelism.unparallelize_pipeline(model, common_device=None)[source]#

Unparallelize a model that was parallelized with parallelize_pipeline. This will remove any pipeline parallelization and move the model to the common device.

Parameters:

  • model (Module) – The model to unparallelize.

  • common_device (Union[device, str, None]) – The device to move the model to. If None, the model will not be moved.

Return type:

Module

Returns:

The model after unparallelization and device transfer.