MXNET教程Fine-Tune训练图像分类模型

1、数据源准备阶段(Data Preparation)

图像分类任务可以算是深度学习的基础也是挑战,著名的数据集包括CIFAR10、CIFAR100、MNIC、Fashion-MNIST等。那么这次使用MXNET框架进行图像分类任务学习练手尝试一下,大家也可以直接参考MXNET官方教程。进行图像分类训练之前,首先要准备你自己的分类数据集。以我学习的训练数据为例进行介绍。数据准备见下图:

见上图,数据集总共有5个类别:分别为cat dog pig duck chicken,其中data目录下面分为train val test三个子文件夹,每个子文件夹下面分别存放5个类别的文件夹。5个类别的文件夹下面存放各自对应的图片数据,记住要检查是否存在图片损坏。文件夹val test存放与图示train文件夹一致

2、预训练模型配置与训练参数设置
2.1 预训练模型下载与路径设置

如果你连接外网的话,直接下载transfer_learning_minc.py文件,然后直接跳转下面一个步骤。

如果你训练环境没有连接网络,建议使用docker环境。那么首先你需要下载预训练模型,本文以ResNet-50为例,你先下载MXNET的关于ImageNet的ResNet-50预训练模型,下载地址为:请点击

如果你已经下载好预训练模型,那么接下来你需要将模型拷贝至①docker环境中~/.mxnet/models/ 路径下面,如果没有这个路径,直接新建这个文件夹路径。如果②不是docker环境,离线情况下直接找到根目录~/.mxnet/models/ 下面,将预训练模型拷贝至文件夹下面。关于如何找到上述路径:

1、启动docker 镜像
2、cd ~
3、ls -a 查看是否存在 .mxnet文件夹 如果存在的话 cd .mxnet/models
4、如果不存在.mxnet文件夹,那么直接mkdir .mxnet文件夹,然后cd .mxnet文件夹目录,继续mkdir models文件夹
5、将ResNet-50预训练模型拷贝至该目录下面 
2.2、训练参数设置和数据增广设置

具体修改部分在代码中体现:

修改的部分主要包括:

① 分类个数classes
②迭代次数epochs
③ 学习率lr和per_device_batch_size
④ transforms.RandomResizedCrop() 函数
⑤ 训练数据加载路径
⑥ 预训练模型名称
⑦ 添加模型保存代码

#!/usr/bin/env python
# -*- coding: UTF-8 -*-

# Transfer Learning with Your Own Image Dataset

################################################################################
# Hyperparameters
# ----------

import mxnet as mx
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:

classes = 5                   # 训练类别个数
epochs = 10000                # 训练全部数据集迭代多少次
lr = 0.001                    # 学习率,因为是Fine-Tune 可以更小
per_device_batch_size = 64    # 每块GPU读取的batch_size
momentum = 0.9
wd = 0.0001

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

num_gpus = 1                  # 建议使用1块GPU训练,我尝试多块GPU不太稳定
num_workers = 8
ctx = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu()]
batch_size = per_device_batch_size * max(num_gpus, 1)

################################################################################
# Things to keep in mind:
# 数据增广部分

jitter_param = 0.4
lighting_param = 0.1

transform_train = transforms.Compose([
    transforms.RandomResizedCrop(128),     # 这里参数设置根据你的图片大小进行设置,设置过大很容易在训练时候抛出Exception
    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(128),                # 这里参数设置根据你的图片大小进行设置,设置过大很容易在训练时候抛出Exception
    #transforms.CenterCrop(224),
    transforms.ToTensor(),
    #transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

################################################################################
# With the data augmentation functions, we can define our data loaders:

path = '$your_path$/data/'               # 训练数据路径,上图中data
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)

################################################################################
#
# Note that only ``train_data`` uses ``transform_train``, while
# ``val_data`` and ``test_data`` use ``transform_test`` to produce deterministic
# results for evaluation.
#
# Model and Trainer
# -----------------
# We use a pre-trained ``ResNet50_v2`` model, which has balanced accuracy and
# computation cost.

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()

################################################################################
# Here's an illustration of the pre-trained model

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()

################################################################################
# Training Loop

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):
        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)

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

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

    _, val_acc = test(finetune_net, val_data, ctx)
    # 添加模型保存代码
    if epoch % 50 == 0:
        finetune_net.export('./minc_models/resnet_model', epoch=epoch)

_, test_acc = test(finetune_net, test_data, ctx)
print('[Finished] Test-acc: %.3f' % (test_acc))
3、相关问题整理
3.1 抛出Exception的异常错误

可能的原因在于你设置 transforms.RandomResizedCrop()函数的参数大小导致的,需要小心对其进行设置。

4、总结

MXNET官方教程还是比较详细,不过在训练图像分类任务进行Fine-Tune时候还是踩了不少坑,主要还是对MXNET框架不是很熟悉。这次的Fine-Tune训练成功,可以改变预训练模型名称,相对应的下载预训练模型就可以进行不同网络的训练进行性能对比。后期有时间的话会进行不同网络的对比,也对之前在数据增广参数设置不确定抛出异常情况要进行研读对应的源码部分。

好了,本文如有错误之处,还请批评指正!

参考连接

https://gluon-cv.mxnet.io/build/examples_classification/transfer_learning_minc.html

猜你喜欢

转载自blog.csdn.net/Small_Munich/article/details/86728674