mxnet.device¶
Device management API of mxnet.
Classes
|
Constructs a device structure. |
Functions
|
Returns a CPU device. |
|
Returns a CPU pinned memory device. |
Returns the current device. |
|
|
Returns a GPU device. |
|
Query CUDA for the free and total bytes of GPU global memory. |
|
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
Returns the device type of current device.
Methods
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
-
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
-
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
-
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