用mxnet实现cnn图像分类

import numpy as np
import os, time, shutil

from mxnet import gluon, image, init, nd
from mxnet import autograd as ag
from mxnet.gluon import nn
from mxnet.gluon.data.vision import transforms
from gluoncv.utils import makedirs
from gluoncv.model_zoo import get_model

We set the hyperparameters as following:

save_frequency = 1

模型保存路径

save_dir = ‘./save_model’

类别数量

classes = 3

全部数据集迭代次数

epochs = 30
lr = 0.001

如果有多个gpu可以用

per_device_batch_size = 2

优化器参数

momentum = 0.9

权值衰减参数，正则作用，一抑制过拟合

wd = 0.0001

学习率调节参数

lr_factor = 0.75
lr_steps = [10, 20, 30, np.inf]

是否使用gpu

num_gpus = 0

是否调用多线程

num_workers = 0
ctx = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu()]

图像数据通过代码加载非常耗内存，所以每轮训练要分batch进行，不然谁的设备也受不了

batch_size = per_device_batch_size * max(num_gpus, 1)

batch_size = 2 # 硬件受限，只能选这么小试试。换言之我们要跑完整个数据集要上百轮，而这仅仅只是一个epoch而已

看一下每张图都是三维的 2562563,非常庞大的数据，远远大于一般的表格化数据

图像增强参数

jitter_param = 0.4
lighting_param = 0.1

transform_train = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomFlipLeftRight(),
transforms.RandomColorJitter(brightness=jitter_param, contrast=jitter_param,
saturation=jitter_param),
transforms.RandomLighting(lighting_param),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

transform_test = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

训练集数据路径

path = ‘./tinyData’
train_path = os.path.join(path, ‘train’)
val_path = os.path.join(path, ‘val’)
test_path = os.path.join(path, ‘test’)

train_data = gluon.data.DataLoader(
gluon.data.vision.ImageFolderDataset(train_path).transform_first(transform_train),
batch_size=batch_size, shuffle=True, num_workers=num_workers)

val_data = gluon.data.DataLoader(
gluon.data.vision.ImageFolderDataset(val_path).transform_first(transform_test),
batch_size=batch_size, shuffle=False, num_workers = num_workers)

test_data = gluon.data.DataLoader(
gluon.data.vision.ImageFolderDataset(test_path).transform_first(transform_test),
batch_size=batch_size, shuffle=False, num_workers = num_workers)

选用模型，

model_name = ‘ResNet50_v2’
finetune_net = get_model(model_name, pretrained=True)

模型参数初始化

with finetune_net.name_scope():
finetune_net.output = nn.Dense(classes)
finetune_net.output.initialize(init.Xavier(), ctx = ctx)
finetune_net.collect_params().reset_ctx(ctx)
finetune_net.hybridize()

训练器，权重参数的收集，优化器的选择，学习率，动量参数等等都在这里设置，这个都是已经封装好的，用就行了

trainer = gluon.Trainer(finetune_net.collect_params(), ‘sgd’, {
‘learning_rate’: lr, ‘momentum’: momentum, ‘wd’: wd})
metric = mx.metric.Accuracy()

定义损失函数，这里用的交叉熵

L = gluon.loss.SoftmaxCrossEntropyLoss()

定义测试函数

def test(net, val_data, ctx):
metric = mx.metric.Accuracy()
for i, batch in enumerate(val_data):
data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0, even_split=False)
label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0, even_split=False)
outputs = [net(X) for X in data]
metric.update(label, outputs)

return metric.get()

lr_counter = 0
num_batch = len(train_data)

开始训练

for epoch in range(epochs):
if epoch == lr_steps[lr_counter]:
trainer.set_learning_rate(trainer.learning_rate*lr_factor)
lr_counter += 1

tic = time.time()
train_loss = 0
metric.reset()

for i, batch in enumerate(train_data):  # 分batch取数据
    data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0, even_split=False)
    label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0, even_split=False)
    with ag.record():
        outputs = [finetune_net(X) for X in data]
        loss = [L(yhat, y) for yhat, y in zip(outputs, label)]
    for l in loss:
        l.backward()

    trainer.step(batch_size)
    train_loss += sum([l.mean().asscalar() for l in loss]) / len(loss)   #loss

    metric.update(label, outputs)

_, train_acc = metric.get()
train_loss /= num_batch

_, val_acc = test(finetune_net, val_data, ctx)

print('[Epoch %d] Train-acc: %.3f, loss: %.3f | Val-acc: %.3f | time: %.1f' %
         (epoch, train_acc, train_loss, val_acc, time.time() - tic))

if save_frequency and save_dir and (epoch + 1) % save_frequency == 0:

    # 模型文件保存，每个save_frequency个epoch保存一次
    finetune_net.save_parameters('%s/imagenet-%s-%d.params' % (save_dir, model_name, epoch))
    trainer.save_states('%s/imagenet-%s-%d.states' % (save_dir, model_name, epoch))

_, test_acc = test(finetune_net, test_data, ctx)
print(’[Finished] Test-acc: %.3f’ % (test_acc))