[MXNet逐梦之旅]练习五·使用MXNetFashionMNIST数据集CNN分类(对比CPU与GPU)
- 使用下述代码实现检测当前环境GPU是否支持
def try_gpu():
try:
ctx = mx.gpu()
_ = mx.nd.zeros((1,), ctx=ctx)
except mx.base.MXNetError:
ctx = mx.cpu()
return ctx
ctx = try_gpu()
- 使用下述代码实现将CPU数据转换到GPU上(运算操作也会自动转换到GPU)如果ctx为CPU则保持在CPU上
for x,y in data_iter:
x, y = x.as_in_context(ctx), y.as_in_context(ctx)#如果ctx为CPU则保持在CPU上
- code
from mxnet import gluon as gl
import mxnet as mx
import time
import sys
import numpy as np
mnist_train = gl.data.vision.FashionMNIST(root="L1/fashion-mnist/",train=True)
mnist_test = gl.data.vision.FashionMNIST(root="L1/fashion-mnist/",train=False)
batch_size = 100
transformer = gl.data.vision.transforms.ToTensor()
if sys.platform.startswith('win'):
num_workers = 0 # 0表示不用额外的进程来加速读取数据
else:
num_workers = 4
train_iter = gl.data.DataLoader(mnist_train.transform_first(transformer),
batch_size, shuffle=True,
num_workers=num_workers)
test_iter = gl.data.DataLoader(mnist_test.transform_first(transformer),
2000, shuffle=False,
num_workers=num_workers)
model = gl.nn.Sequential()
model.add(gl.nn.Conv2D(16,5,activation="sigmoid"),
gl.nn.MaxPool2D(2,2),
gl.nn.Conv2D(64,5,activation="sigmoid"),
gl.nn.MaxPool2D(2,2),
gl.nn.Dense(128,activation="relu"),
gl.nn.Dense(10))
def try_gpu():
try:
ctx = mx.gpu()
_ = mx.nd.zeros((1,), ctx=ctx)
except mx.base.MXNetError:
ctx = mx.cpu()
return ctx
ctx = try_gpu()
model.initialize(force_reinit=True, ctx=ctx, init=mx.init.Xavier())
loss = gl.loss.SoftmaxCrossEntropyLoss()
lr = 0.001
opt = gl.Trainer(model.collect_params(),"adam",{"learning_rate":lr})
def accuracy(y_hat, y):
return (y_hat.argmax(axis=1) == y.astype('float32')).mean().asscalar()
def train(data_iter,model,loss,opt,ctx,Num,batch_size):
print("\nTrain on",ctx)
for e in range(1,Num+1):
ts = time.time()
losses = []
for x,y in data_iter:
x, y = x.as_in_context(ctx), y.as_in_context(ctx)
with mx.autograd.record():
l = loss(model(x),y)
losses.append(l.asnumpy())
l.backward()
opt.step(batch_size)
loss_val = np.mean(losses)
accy = accuracy(model(x),y)
print("Epoch %d, loss: %f acc: %f, time: %f"%(e,loss_val,accy, time.time()-ts))
train(train_iter,model,loss,opt,ctx,5,batch_size)
- out
GPU:
Train on gpu(0)
Epoch 1, loss: 1.051915 acc: 0.700000, time: 6.906520
Epoch 2, loss: 0.581923 acc: 0.750000, time: 6.714078
Epoch 3, loss: 0.477215 acc: 0.820000, time: 6.866763
Epoch 4, loss: 0.416427 acc: 0.860000, time: 6.939581
Epoch 5, loss: 0.382869 acc: 0.850000, time: 6.999164
CPU:
Train on cpu(0)
Epoch 1, loss: 1.047916 acc: 0.770000, time: 57.309021
Epoch 2, loss: 0.580764 acc: 0.860000, time: 60.355942
Epoch 3, loss: 0.475856 acc: 0.840000, time: 61.056674
Epoch 4, loss: 0.415988 acc: 0.890000, time: 60.469341
Epoch 5, loss: 0.385315 acc: 0.880000, time: 60.729165