MXNet使用GPU计算

GPU上的存储

在MXNet中，mx.cpu()代表所有的物理CPU和内存，而mx.gpu()只代表一块GPU和相应的显存。(cpu对应内存，gpu对应显存)

可以通过NDArray的context属性来查看该NDArray所在的设备。

x.context

我们可以通过多种方法创建NDArray到GPU的显存上。

a = nd.array([1,2,3], ctx=mx.gpu()) # gpu()等价于gpu(0)

# 假设有2块GPU的情况下
B = nd.random.uniform(shape=(2, 3), ctx=mx.gpu(1))

通过copyto函数和as_in_context函数在设备之间传输数据。

y = x.copyto(mx.gpu())

z = x.as_in_context(mx.gpu())

需要注意的是，如果x和y在同一个设备上，as_in_context函数不会开辟新内存，而是让目标变量共享原变量的内存或或者显存。

而copyto函数总是为目标变量开辟新的内存或者显存。