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函数总是为目标变量开辟新的内存或者显存。
GPU上的计算
MXNet的计算会在数据的context属性所指定的设备上执行。并且计算结果自动保存在同一块显卡的显存上。
MXNet要求计算的所有输入数据都应该在同一个设备上,即都在内存上或都在同一块显卡的显存上。
和NDArray类似,Gluon模型可以在初始化时通过ctx参数指定设备,使模型参数初始化在相应的设备上。
net = nn.Sequential()
net.add(nn.Dense(1))
net.initialize(ctx=mx.gpu())