mxnet.context¶

Context management API of mxnet.

Classes

Functions

class mxnet.context.Context(device_type, device_id=0)[source]¶

Bases: object

Constructs a context.

MXNet can run operations on CPU and different GPUs. A context 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 context.

Methods

empty_cache()

Empties the memory cache for the current contexts device.

Parameters

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

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

Note

Context can also be used as a way to change the default context.

Examples

>>> # array on cpu
>>> cpu_array = mx.nd.ones((2, 3))
>>> # switch default context to GPU(2)
>>> with mx.Context(mx.gpu(2)):
...     gpu_array = mx.nd.ones((2, 3))
>>> gpu_array.context
gpu(2)

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

>>> gpu_array = mx.nd.ones((2, 3), mx.gpu(1))
>>> gpu_array.context
gpu(1)

property device_type¶

Returns the device type of current context.

Examples

>>> mx.context.current_context().device_type
'cpu'
>>> mx.current_context().device_type
'cpu'

Returns

device_type

Return type

str

empty_cache()[source]¶

Empties the memory cache for the current contexts device.

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

Examples

>>> ctx = mx.gpu(0)
>>> arr = mx.nd.ones((200,200), ctx=ctx)
>>> del arr
>>> ctx.empty_cache() # forces release of memory allocated for arr

mxnet.context.cpu(device_id=0)[source]¶

Returns a CPU context.

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

Examples

>>> with mx.cpu():
...     cpu_array = mx.nd.ones((2, 3))
>>> cpu_array.context
cpu(0)
>>> cpu_array = mx.nd.ones((2, 3), ctx=mx.cpu())
>>> cpu_array.context
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

context – The corresponding CPU context.

Return type

Context

mxnet.context.cpu_pinned(device_id=0)[source]¶

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

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

Examples

>>> with mx.cpu_pinned():
...     cpu_array = mx.nd.ones((2, 3))
>>> cpu_array.context
cpu_pinned(0)
>>> cpu_array = mx.nd.ones((2, 3), ctx=mx.cpu_pinned())
>>> cpu_array.context
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

context – The corresponding CPU pinned memory context.

Return type

Context

mxnet.context.current_context()[source]¶

Returns the current context.

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

Examples

>>> mx.current_context()
cpu(0)
>>> with mx.Context('gpu', 1):  # Context changed in `with` block.
...    mx.current_context()  # Computation done here will be on gpu(1).
...
gpu(1)
>>> mx.current_context() # Back to default context.
cpu(0)

Returns

default_ctx

Return type

Context

mxnet.context.gpu(device_id=0)[source]¶

Returns a GPU context.

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

Examples

>>> cpu_array = mx.nd.ones((2, 3))
>>> cpu_array.context
cpu(0)
>>> with mx.gpu(1):
...     gpu_array = mx.nd.ones((2, 3))
>>> gpu_array.context
gpu(1)
>>> gpu_array = mx.nd.ones((2, 3), ctx=mx.gpu(1))
>>> gpu_array.context
gpu(1)

Parameters

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

Returns

context – The corresponding GPU context.

Return type

Context

mxnet.context.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)

mxnet.context.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