mxnet.device

Device management API of mxnet.

Classes

Device(device_type[, device_id])

Constructs a device structure.

Functions

cpu([device_id])

Returns a CPU device.

cpu_pinned([device_id])

Returns a CPU pinned memory device.

current_device()

Returns the current device.

gpu([device_id])

Returns a GPU device.

gpu_memory_info([device_id])

Query CUDA for the free and total bytes of GPU global memory.

num_gpus()

Query CUDA for the number of GPUs present.

class Device(device_type, device_id=0)[source]

Bases: object

Constructs a device structure.

MXNet can run operations on CPU and different GPUs. A Device class describes the device type and ID on which computation should be carried on.

One can use mx.cpu and mx.gpu for short.

See also

How to run MXNet on multiple CPU/GPUs <http://mxnet.incubator.apache.org/api/faq/distributed_training> for more details.

Attributes

device_type

Returns the device type of current device.

Methods

empty_cache()

Empties the memory cache for the current device.

Parameters
  • device_type ({'cpu', 'gpu'} or Device.) – String representing the device type.

  • device_id (int (default=0)) – The device id of the device, needed for GPU.

Note

Device can also be used as a way to change the default device.

Examples

>>> # array on cpu
>>> cpu_array = mx.np.ones((2, 3))
>>> # switch default Device to GPU(2)
>>> with mx.Device(mx.gpu(2)):
...     gpu_array = mx.np.ones((2, 3))
>>> gpu_array.device
gpu(2)

One can also explicitly specify the device when creating an array.

>>> gpu_array = mx.np.ones((2, 3), mx.gpu(1))
>>> gpu_array.device
gpu(1)
property device_type

Returns the device type of current device.

Examples

>>> mx.device.current_device().device_type
'cpu'
>>> mx.current_device().device_type
'cpu'
Returns

device_type

Return type

str

empty_cache()[source]

Empties the memory cache for the current device.

MXNet utilizes a memory pool to avoid excessive allocations. Calling empty_cache will empty the memory pool of the device. This will only free the memory of the unreferenced data.

Examples

>>> ctx = mx.gpu(0)
>>> arr = mx.np.ones((200,200), ctx=ctx)
>>> del arr
>>> ctx.empty_cache() # forces release of memory allocated for arr
cpu(device_id=0)[source]

Returns a CPU device.

This function is a short cut for Device('cpu', device_id). For most operations, when no device is specified, the default device is cpu().

Examples

>>> with mx.cpu():
...     cpu_array = mx.np.ones((2, 3))
>>> cpu_array.device
cpu(0)
>>> cpu_array = mx.np.ones((2, 3), ctx=mx.cpu())
>>> cpu_array.device
cpu(0)
Parameters

device_id (int, optional) – The device id of the device. device_id is not needed for CPU. This is included to make interface compatible with GPU.

Returns

device – The corresponding CPU device.

Return type

Device

cpu_pinned(device_id=0)[source]

Returns a CPU pinned memory device. Copying from CPU pinned memory to GPU is faster than from normal CPU memory.

This function is a short cut for Device('cpu_pinned', device_id).

Examples

>>> with mx.cpu_pinned():
...     cpu_array = mx.np.ones((2, 3))
>>> cpu_array.device
cpu_pinned(0)
>>> cpu_array = mx.np.ones((2, 3), ctx=mx.cpu_pinned())
>>> cpu_array.device
cpu_pinned(0)
Parameters

device_id (int, optional) – The device id of the device. device_id is not needed for CPU. This is included to make interface compatible with GPU.

Returns

device – The corresponding CPU pinned memory device.

Return type

Device

current_device()[source]

Returns the current device.

By default, mx.cpu() is used for all the computations and it can be overridden by using with mx.Device(x) statement where x can be cpu(device_id) or gpu(device_id).

Examples

>>> mx.current_device()
cpu(0)
>>> with mx.Device('gpu', 1):  # Device changed in `with` block.
...    mx.current_device()  # Computation done here will be on gpu(1).
...
gpu(1)
>>> mx.current_device() # Back to default device.
cpu(0)
Returns

default_device

Return type

Device

gpu(device_id=0)[source]

Returns a GPU device.

This function is a short cut for Device(‘gpu’, device_id). The K GPUs on a node are typically numbered as 0,…,K-1.

Examples

>>> cpu_array = mx.np.ones((2, 3))
>>> cpu_array.device
cpu(0)
>>> with mx.gpu(1):
...     gpu_array = mx.np.ones((2, 3))
>>> gpu_array.device
gpu(1)
>>> gpu_array = mx.np.ones((2, 3), ctx=mx.gpu(1))
>>> gpu_array.device
gpu(1)
Parameters

device_id (int, optional) – The device id of the device, needed for GPU.

Returns

device – The corresponding GPU device.

Return type

Device

gpu_memory_info(device_id=0)[source]

Query CUDA for the free and total bytes of GPU global memory.

Parameters

device_id (int, optional) – The device id of the GPU device.

Raises

Will raise an exception on any CUDA error.

Returns

(free, total)

Return type

(int, int)

num_gpus()[source]

Query CUDA for the number of GPUs present.

Raises

Will raise an exception on any CUDA error.

Returns

count – The number of GPUs.

Return type

int