mednet.engine.device¶
Support for switching execution devices (GPU vs CPU).
Module Attributes
List of supported pytorch devices by this library. |
Classes
|
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
, orcuda: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
, orcuda:3
). In the specific case ofcuda
, 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: