jacinle.cli.device#

Uiltity functions to parse CPU/CUDA device strings.

Classes

DeviceNameFormat

The target fmt of device names.

Functions

`canonlize_device_name`(d[, fmt])	Convert a device name to a canonical format.
`parse_and_set_devices`(devs[, fmt, set_device])	Parse a input list of strings or a single comma-separated string into a list of device names.
`parse_devices`(devs[, fmt])	Parse a input list of strings or a single comma-separated string into a list of device names.
`set_cuda_visible_devices`(devs)	Set the CUDA_VISIBLE_DEVICES environment variable with a single comma-separated string or a list of strings.

Class DeviceNameFormat

class DeviceNameFormat[source]#

Bases: JacEnum

The target fmt of device names. Supported formats are:

DeviceNameFormat.INT: integer, e.g., 0, 1, 2, etc.
DeviceNameFormat.TENSORFLOW: TensorFlow-style device name, e.g., /cpu:0, /gpu:1, etc.

__new__(value)#

classmethod assert_valid(value)#: Assert if the value is a valid choice.

classmethod choice_names()#: Returns the list of the name of all possible choices.

classmethod choice_objs()#: Returns the list of the object of all possible choices.

classmethod choice_values()#: Returns the list of the value of all possible choices.

classmethod from_string(value)#

Parameters:: value (str | JacEnum)
Return type:: JacEnum

classmethod is_valid(value)#: Check if the value is a valid choice.

classmethod type_name()#: Return the type name of the enum.

INT = 'int'#

TENSORFLOW = 'tensorflow'#

Functions

canonlize_device_name(d, fmt=DeviceNameFormat.INT)[source]#

Convert a device name to a canonical format.

Parameters:

d (str) – the device name to be converted. The string can be either: cpu, gpu0, or 0.
fmt (str | DeviceNameFormat) – the target format.

Returns:

the canonical device name. If the target format is DeviceNameFormat.INT, the return value is an integer. When d is cpu, the return value is -1. If the target format is DeviceNameFormat.TENSORFLOW, the return value is a string: e.g., /cpu:0, /gpu:1, etc.

Return type:

str | int

parse_and_set_devices(devs, fmt=DeviceNameFormat.INT, set_device=True)[source]#

Parse a input list of strings or a single comma-separated string into a list of device names. When set_device is True, the CUDA_VISIBLE_DEVICES environment variable will be set accordingly.

Parameters:

devs (str | Sequence[str]) – the input device list.
fmt (str | DeviceNameFormat) – the target format.
set_device (bool) – whether to set the CUDA_VISIBLE_DEVICES environment variable.

Returns:

the parsed device list.

parse_devices(devs, fmt=DeviceNameFormat.INT)[source]#

Parse a input list of strings or a single comma-separated string into a list of device names.

Parameters:

devs (str | Sequence[str]) – the input device list.
fmt (str | DeviceNameFormat) – the target format.

Returns:

the parsed device list.

Return type:

List[str | int]

set_cuda_visible_devices(devs)[source]#

Set the CUDA_VISIBLE_DEVICES environment variable with a single comma-separated string or a list of strings.

Parameters:: devs (str | Sequence[str]) – the input device list. Either a single comma-separated string or a list of strings.