[MXNet逐梦之旅]练习五·使用MXNetFashionMNIST数据集CNN分类（对比CPU与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()

for x,y in data_iter:
    x, y = x.as_in_context(ctx), y.as_in_context(ctx)#如果ctx为CPU则保持在CPU上

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)

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

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