MXnet-gluoncv实现基于迁移学习的目标检测（自己的数据集）

一、安装

首先安装mxnet（cpu版）：
命令行输入pip install -i https://pypi.doubanio.com/simple/ mxnet（豆瓣源、速度快）
接着安装gluoncv（cpu版）：
命令行输入pip install --upgrade mxnet gluoncv（如果直接安装gluoncv，会从国外源下载mxnet安装包，速度慢到报错……建议还是先装mxnet再装gluoncv吧）

二、数据集制作

先自制PASCAL VOC格式的数据集，再将PASCAL VOC格式的数据集转换为适用于MXnet框架的数据集。

from gluoncv.data import VOCDetection
class VOCLike(VOCDetection):CLASSES = ['xxx1', 'xxx2']  #自己数据集中的类别def __init__(self, root, splits, transform=None, index_map=None, preload_label=True):super(VOCLike, self).__init__(root, splits, transform, index_map, preload_label)dataset = VOCLike(root='./VOCtemplate', splits=((2007, 'traineval'),))      #自己数据集的路径，包含VOC2007的上一层目录
print('length of dataset:', len(dataset))
print('label example:')
print(dataset[0][1])

报错：
d:\Anaconda3\lib\site-packages\gluoncv\data\pascal_voc\detection.py in _load_label(self, idx)
109 for obj in root.iter(‘object’):
110 try:
–> 111 difficult = int(obj.find(‘difficult’).text)
112 except ValueError:
113 difficult = 0

AttributeError: ‘NoneType’ object has no attribute ‘text’

解决：
我数据集的标签是用labelimg做的，标签中的difficult都写作“Difficult”，因此把gluoncv\data\pascal_voc\detection.py中的‘difficult’改为‘Difficult’即可。

三、迁移学习—以ssd_512_mobilenet1.0为例

#加载预训练过的模型
net = gcv.model_zoo.get_model('ssd_512_mobilenet1.0_voc', pretrained=True)
#重设类别
classes=['xxx1', 'xxx2']  #自己数据集中的类别
net.reset_class(classes)
#batch设置，规范输入网络进行训练的数据格式
def get_dataloader(net, train_dataset, data_shape, batch_size, num_workers):from gluoncv.data.batchify import Tuple, Stack, Padfrom gluoncv.data.transforms.presets.ssd import SSDDefaultTrainTransformwidth, height = data_shape, data_shape# use fake data to generate fixed anchors for target generationwith autograd.train_mode():_, _, anchors = net(mx.nd.zeros((1, 3, height, width)))batchify_fn = Tuple(Stack(), Stack(), Stack())  # stack image, cls_targets, box_targetstrain_loader = gluon.data.DataLoader(train_dataset.transform(SSDDefaultTrainTransform(width, height, anchors)),batch_size, True, batchify_fn=batchify_fn, last_batch='keep', num_workers=num_workers)return train_loadertrain_data = get_dataloader(net, dataset, 512, 4, 0)

设置用于训练的处理器类型

#尝试用gpu训练
try:a = mx.nd.zeros((1,), ctx=mx.gpu(0))ctx = [mx.gpu(0)]
except:ctx = [mx.cpu()]

开始训练：

net.collect_params().reset_ctx(ctx)
trainer = gluon.Trainer(net.collect_params(), 'sgd',{'learning_rate': 0.001, 'wd': 0.0005, 'momentum': 0.9})mbox_loss = gcv.loss.SSDMultiBoxLoss()
ce_metric = mx.metric.Loss('CrossEntropy')
smoothl1_metric = mx.metric.Loss('SmoothL1')for epoch in range(0, 10):    #设置epoch数ce_metric.reset()smoothl1_metric.reset()tic = time.time()btic = time.time()net.hybridize(static_alloc=True, static_shape=True)for i, batch in enumerate(train_data):batch_size = batch[0].shape[0]data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)cls_targets = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0)box_targets = gluon.utils.split_and_load(batch[2], ctx_list=ctx, batch_axis=0)with autograd.record():cls_preds = []box_preds = []for x in data:cls_pred, box_pred, _ = net(x)cls_preds.append(cls_pred)box_preds.append(box_pred)sum_loss, cls_loss, box_loss = mbox_loss(cls_preds, box_preds, cls_targets, box_targets)autograd.backward(sum_loss)# since we have already normalized the loss, we don't want to normalize# by batch-size anymoretrainer.step(1)ce_metric.update(0, [l * batch_size for l in cls_loss])smoothl1_metric.update(0, [l * batch_size for l in box_loss])name1, loss1 = ce_metric.get()name2, loss2 = smoothl1_metric.get()if i % 20 == 0:print('[Epoch {}][Batch {}], Speed: {:.3f} samples/sec, {}={:.3f}, {}={:.3f}'.format(epoch, i, batch_size/(time.time()-btic), name1, loss1, name2, loss2))btic = time.time()

保存训练好的模型参数：

net.save_parameters('./ssd_512_mobilenet1.0_mound.params') #设置模型参数的存放路径

四、测试


net = gcv.model_zoo.get_model('ssd_512_mobilenet1.0_custom', classes=classes, pretrained_base=False)
net.load_parameters('./ssd_512_mobilenet1.0_mound.params') #载入模型参数
x, image = gcv.data.transforms.presets.ssd.load_test('./xxxx.jpg', 512) #读取测试图像
cid, score, bbox = net(x)
ax = viz.plot_bbox(image, bbox[0], score[0], cid[0], class_names=classes)
plt.show()

代码参考gluoncv的 tutorial
https://gluon-cv.mxnet.io/build/examples_detection/finetune_detection.html