mxnet GPU加速计算

 

#coding:utf-8
'''
NDArray在GPU上计算
'''

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

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

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

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

print('a = ',a)
print('b = ',b)
print('x = ',x)
print('y = ',y)
print('z = ',z)

# Gluon的GPU计算

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

data = nd.random.uniform(shape=(3,2),ctx=mx.gpu())
print(net(data))

# 模型的参数也储存在GPU上
print(net[0].weight.data())