NDArray Contexts¶
Overview¶
This guide will introduce you to managing CPU versus GPU contexts for handling data.
This content was extracted and simplified from the gluon tutorials in Dive Into Deep Learning.
Prerequisites¶
MXNet installed (with GPU support) in a Python environment.
Python 2.7.x or Python 3.x
One or more GPUs
Managing Context¶
In MXNet, every array has a context. One context could. If you have a GPU, let’s try initializing an array on the first GPU.
Otherwise, use
ctx=mx.cpu() in place of
ctx=gpu(0).
[ ]:
from mxnet import gpu from mxnet import nd z = nd.ones(shape=(3,3), ctx=gpu(0)) print(z)
Given an NDArray on a given context, we can copy it to another context by using the copyto() method. Skip this if you don’t have a GPU at the moment.
[ ]:
x_gpu = x.copyto(gpu(0)) print(x_gpu)
The result of an operator will have the same context as the inputs.
[ ]:
x_gpu + z
Watch out!¶
Imagine that your variable z already lives on your second GPU (
gpu(0)). What happens if we call
z.copyto(gpu(0))? It will make a copy and allocate new memory, even though that variable already lives on the desired device!
Often, we only want to make a copy if the variable currently lives in the wrong context. In these cases, we can call
as_in_context(). If the variable is already on
gpu(0) then this is a no-op.
[ ]:
print('id(z):', id(z)) z = z.copyto(gpu(0)) print('id(z):', id(z)) z = z.as_in_context(gpu(0)) print('id(z):', id(z)) print(z)
Next Up¶
Back to NDArray API Guides