【MXNet】(十八):使用GPU计算

版权声明:本文为博主原创文章,未经作者允许请勿转载。 https://blog.csdn.net/heiheiya https://blog.csdn.net/heiheiya/article/details/87259635

MXNet可以指定用来计算和存储的设备。默认情况下,MXNet会将数据创建在内存,然后利用CPU来计算。在MXNet中,mx.cpu()(或者在括号里填任意整数)表示所有的物理CPU和内存。这意味着,MXNet的计算会尽量使用所有的CPU核。但mx.gpu()只代表一块GPU和相应的显存。如果有多块GPU,我们用mx.gpu(i)来表示第ii块GPU及相应的显存(ii从0开始)且mx.gpu(0)mx.gpu()等价。

import mxnet as mx
from mxnet import nd
from mxnet.gluon import nn

mx.cpu(), mx.gpu(), mx.gpu(0)

在默认情况下,NDArray存在内存上。

x = nd.array([1, 2, 3])
x

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

x.context

可以在创建NDArray的时候通过ctx参数指定存储设备。

a = nd.array([1, 2, 3], ctx=mx.gpu())
a

B = nd.random.uniform(shape=(2, 3), ctx=mx.gpu(0))
B

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

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

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

如果源变量和目标变量的context一致,as_in_context函数使目标变量和源变量共享源变量的内存或显存。

y.as_in_context(mx.gpu()) is y

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

y.copyto(mx.gpu()) is y

MXNet的计算会在数据的context属性所指定的设备上执行。为了使用GPU计算,我们只需要事先将数据存储在显存上。计算结果会自动保存在同一块显卡的显存上。MXNet要求计算的所有输入数据都在内存或同一块显卡的显存上。

(z + 2).exp() * y

Gluon的模型可以在初始化时通过ctx参数指定设备。

net = nn.Sequential()
net.add(nn.Dense(1))
net.initialize(ctx=mx.gpu())
net(y)

看一下模型参数是不是存储在同一块显卡的显存上。

net[0].weight.data()

展开阅读全文

没有更多推荐了,返回首页