mxnet.context¶

Context management API of mxnet.

Classes

Context(device_type[, device_id])

Constructs a context.

Functions

`cpu`([device_id])	Returns a CPU context.
`cpu_pinned`([device_id])	Returns a CPU pinned memory context.
`current_context`()	Returns the current context.
`gpu`([device_id])	Returns a GPU context.
`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 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

Did this page help you?

Yes

No

Thanks for your feedback!