MXNet tutorial——Train the neural network

最新推荐文章于 2024-08-18 12:47:57 发布

xw2017

最新推荐文章于 2024-08-18 12:47:57 发布

阅读量307

点赞数

分类专栏： MXNet 文章标签： mxnet

MXNet 专栏收录该内容

1 篇文章 0 订阅

订阅专栏

数据集：FashionMNIST

导入类库

from mxnet import nd, gluon, init, autograd
from mxnet.gluon import nn
from mxnet.gluon.data.vision import datasets, transforms
import time

准备数据

// 训练集
mnist_train = datasets.FashionMNIST(train=True)
transformer = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(0.13, 0.31)])
mnist_train = mnist_train.transform_first(transformer)
// 验证集
mnist_valid = gluon.data.vision.FashionMNIST(train=False)
valid_data = gluon.data.DataLoader(
    mnist_valid.transform_first(transformer),
    batch_size=batch_size, num_workers=4)

定义模型

// 定义模型
net = nn.Sequential()
net.add(nn.Conv2D(channels=6, kernel_size=5, activation='relu'),
        nn.MaxPool2D(pool_size=2, strides=2),
        nn.Conv2D(channels=16, kernel_size=3, activation='relu'),
        nn.MaxPool2D(pool_size=2, strides=2),
        nn.Flatten(),
        nn.Dense(120, activation="relu"),
        nn.Dense(84, activation="relu"),
        nn.Dense(10))
// 参数初始化
net.initialize(init=init.Xavier())
// 定义损失函数(交叉熵损失函数)
softmax_cross_entropy = gluon.loss.SoftmaxCrossEntropyLoss()
// 定义优化器(sgd优化)
trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': 0.1})

训练模型

def acc(output, label):
    return (output.argmax(axis=1) ==
            label.astype('float32')).mean().asscalar()
  
for epoch in range(10):
    train_loss, train_acc, valid_acc = 0., 0., 0.
    tic = time.time()
    for data, label in train_data:
        # 前向传播+反向传播
        with autograd.record():
            output = net(data)
            loss = softmax_cross_entropy(output, label)
        loss.backward()
        # 更新参数
        trainer.step(batch_size)
        # 计算训练精度
        train_loss += loss.mean().asscalar()
        train_acc += acc(output, label)
    # 计算验证精度
    for data, label in valid_data:
        valid_acc += acc(net(data), label)
    print("Epoch %d: loss %.3f, train acc %.3f, test acc %.3f, in %.1f sec" % (
            epoch, train_loss/len(train_data), train_acc/len(train_data),
            valid_acc/len(valid_data), time.time()-tic))