mednet.engine.device

Support for switching execution devices (GPU vs CPU).

Module Attributes

SupportedPytorchDevice

List of supported pytorch devices by this library.

Classes

DeviceManager(name)

Manage Lightning Accelerator and Pytorch Devices.

mednet.engine.device.SupportedPytorchDevice

List of supported pytorch devices by this library.

alias of Literal[‘cpu’, ‘cuda’, ‘mps’]

class mednet.engine.device.DeviceManager(name)[source]

Bases: object

Manage Lightning Accelerator and Pytorch Devices.

It takes the user input, in the form of a string defined by [\S+][:\d[,\d]?]? (e.g.: cpu, mps, or cuda:3), and can translate to the right incarnation of Pytorch devices or Lightning Accelerators to interface with the various frameworks.

Instances of this class also manage the environment variable $CUDA_VISIBLE_DEVICES if necessary.

Parameters:

name (Literal['cpu', 'cuda', 'mps']) – The name of the device to use, in the form of a string defined by [\S+][:\d[,\d]?]? (e.g.: cpu, mps, or cuda:3). In the specific case of cuda, one can also specify a device to use either by adding :N, where N is the zero-indexed board number on the computer, or by setting the environment variable $CUDA_VISIBLE_DEVICES with the devices that are usable by the current process.

torch_device()[source]

Return a representation of the torch device to use by default.

Warning

If a list of devices is set, then this method only returns the first device. This may impact Nvidia GPU logging in the case multiple GPU cards are used.

Returns:

The first torch device (if a list of ids is set).

Return type:

torch.device

lightning_accelerator()[source]

Return the lightning accelerator setup.

Return type:

tuple[str, int | list[int] | str]

Returns:

  • accelerator – The lightning accelerator to use.

  • devices – The lightning devices to use.