cifar10图像分类

 

import argparse #行参数管理
import mxnet as mx
import os
import logging
# os.environ['CUDA_VISIBLE_DEVICES'] = '2'

"""
    网络结构的搭建：主要就是替换原网络结构中的全连接层的操作
    传入的变量有：
    1、预训练模型的网络结构sym
    2、当前分类任务的类别数num_classes
    3、全连接层前面一层的名称
"""
def get_fine_tune_model(sym, num_classes, layer_name):

    #调用所有层的信息
    all_layers = sym.get_internals()

    #在指定层的名称后面加上"_output";然后在指定层后面添加新的全连接层，全连接层的输出节点数与类别数相同
    net = all_layers[layer_name + '_output']
    net = mx.symbol.FullyConnected(data=net, num_hidden=num_classes,
                                   name='new_fc')
    net = mx.symbol.SoftmaxOutput(data=net, name='softmax')
    return net

"""
    该函数可以得到可用的学习率变化对象lr_scheduler
"""
def multi_factor_scheduler(args, epoch_size):

    #这里是计算每多少步更新一次学习率
    step = range(args.step, args.num_epoch, args.step)
    step_bs = [epoch_size * (x - args.begin_epoch) for x in step
             if x - args.begin_epoch > 0]
    if step_bs:
        return mx.lr_scheduler.MultiFactorScheduler(step=step_bs,
                                                    factor=args.factor)
    return None

def data_loader(args):
    """
        1、通过获取预先设置的图像形状“3，224，224”
        图片的形状设置的是：3,224,224 这里通过split(",")进行获取
        然后将其转换成网络所需要的元组类型（3，224，224）

    """
    data_shape_list = [int(item) for item in args.image_shape.split(",")]
    data_shape = tuple(data_shape_list)

    """
        2、mx.io.ImageRecordIter接口实现读取训练数据集的RecordIO文件；同样也可以使用原始的图像文件；但是比较费时费力
    """
    train = mx.io.ImageRecordIter(
        path_imgrec=args.data_train_rec,
        path_imgidx=args.data_train_idx,
        label_width=1, #单标签分类，有的还是多标签的分类；根据实际情况来吧

        #分别对应RGB三个通道的均值，将读入的图像，依次减去输入图像的各通道像素值对应的三个均值
        mean_r=123.68,
        mean_g=116.779,
        mean_b=103.939,

        #与网络构建时的数据和标签相同
        data_name='data',
        label_name='softmax_label',
        data_shape=data_shape,
        batch_size=args.batch_size,

        #对输入的图像进行随机镜像，可以增加数据的多样性
        rand_mirror=args.random_mirror,

        #随机对比度调整，也是数据增强的一种方法，默认为0.3
        max_random_contrast=args.max_random_contrast,

        #随机旋转图像时的最大角度，默认为15，也就是在正负15之间随机旋转输入图像
        max_rotate_angle=args.max_rotate_angle,

        #是否对图像进行随机打乱；训练集可以，但是验证集可以不用这部操作了
        shuffle=True,

        #修改训练数据的尺寸；对于训练集来说，为了增加数据的多样性，通常将尺寸设置为256；对于验证集要尽可能的保存图像的内容了
        resize=args.resize_train)

    """
        2、和训练数据集一样，这里是读取验证集数据，参数的定义也是一样的
    """
    val = mx.io.ImageRecordIter(
        path_imgrec=args.data_val_rec,
        path_imgidx=args.data_val_idx,
        label_width=1,
        mean_r=123.68,
        mean_g=116.779,
        mean_b=103.939,
        data_name='data',
        label_name='softmax_label',
        data_shape=data_shape,
        batch_size=args.batch_size,
        rand_mirror=0,
        shuffle=False,
        resize=args.resize_val)
    return train,val

def train_model(args):
    """
        1、读取训练和验证数据集
    """
    train, val = data_loader(args)

    """
        2、导入预训练模型并调用get_fine_tune_model函数对导入的网络模型进行修改,生成一个新的模型
        load_checkpoint(model_prefix,epoch)
    """
    sym, arg_params, aux_params = mx.model.load_checkpoint(prefix=args.model,
                                                           epoch=args.begin_epoch)
    new_sym = get_fine_tune_model(sym, args.num_classes, args.layer_name)

    """
        3、每隔step个epoch就将学习率呈上一个factor作为一个新的学习率；下面就是将总的数据量除以batch_size
        
    """
    epoch_size = max(int(args.num_examples / args.batch_size), 1)
    lr_scheduler = multi_factor_scheduler(args, epoch_size)

    '''
        4、构造优化相关的字典:
        学习率、
        动量参数可用于随机梯度下降常用0.9、
        正则项参数：权重衰减来防止过拟合、
        以及学习率变化策略：就是可能乘以decay的那个factor
    '''
    optimizer_params = {'learning_rate': args.lr,
                        'momentum': args.mom,
                        'wd': args.wd,
                        'lr_scheduler': lr_scheduler}

    """
        5、设置网络的初始化方式：
        采用高斯函数进行随机初始化
    """
    initializer = mx.init.Xavier(rnd_type='gaussian',
                                 factor_type="in",
                                 magnitude=2)

    """
	    6、用于设置模型训练的环境，默认为0代表在GPU上运行
	    #添加0既可以用
	"""
    if args.gpus == '':
        devs = mx.cpu()
    else:
		#devs = mx.gpu(args.gpus)
        devs = [mx.gpu(int(i)) for i in args.gpus.split(',')]

    """
        7、设置在训练过程中是否固定部分网络层参数，默认为false就是不固定；
        后面会讲到固定预训练层的模型参数，只更新全连接层的参数信息
    """
    if args.fix_pretrain_param:
        fixed_param_names = [layer_name for layer_name in
                             new_sym.list_arguments() if layer_name not in
                             ['new_fc_weight', 'new_fc_bias', 'data',
                              'softmax_label']]
    else:
        fixed_param_names = None

    """
        8、根据前面得到的模型和训练环境初始化得到一个module对象，用于后续的模型训练
    """
    model = mx.mod.Module(context=devs,
                          symbol=new_sym,
                          fixed_param_names=fixed_param_names)

    """
        9、设置日志保存和显示的批次间隔batch_callback,保存的epoch间隔默认是1，简单说就是显示训练到第几个epoch了，第几个batch了，以及训练速度等信息；
            还有训练得到的模型的保存地址epoch_callback；保存的地址加保存的名称
    """
    batch_callback = mx.callback.Speedometer(args.batch_size, args.period)
    epoch_callback = mx.callback.do_checkpoint(args.save_result + args.save_name)

    """
        10、判断是否使用预训练模型参数初始化网络结构
        默认为False,就是使用预先训练的模型参数
        True的话就是使用指定的随机初始化方式
    """
    if args.from_scratch:
        arg_params = None
        aux_params = None

    """
        11、model的fit()方法
        传入指定参数启动训练，这里指定的是准确率acc 和 交叉熵 ce(cross entropy)
    
    """
    model.fit(train_data=train,
              eval_data=val,
              begin_epoch=args.begin_epoch,
              num_epoch=args.num_epoch,
              eval_metric=['acc','ce'],
              optimizer='sgd',
              optimizer_params=optimizer_params,
              arg_params=arg_params,
              aux_params=aux_params,
              initializer=initializer,
              allow_missing=True,
              batch_end_callback=batch_callback,
              epoch_end_callback=epoch_callback)

#命令行参数信息
def parse_arguments():
    """
        argparse命令行参数解析包
        1、其中ArgumentParser是用来生成一个parser对象，然后其中的descriptor用于描述这个解析器是做什么的！
        2、通过parser的add_argument来增加参数；如果不设置参数，也可以自己default,程序会自动将此值当作参数值
        3、其中namespace中的两个属性，当“-”和“--”同时出现时，系统默认后者为参数名；但是在命令行输入的时候则没有这个区分
        4、可以通过python 程序名.py -h来查看帮助信息
        5、num 为一个位置参数、type表示参数的类型、参数默认都是string、help表示参数的描述信息
        6、--mode以--开头，为一个可选参数、可以有多个别名、可选参数不是必选的，一般用来做条件选择的
        7、最终如何获取参数的信息呢？很简单了，使用args加.进行读取里面的参数信息，并且可以很方便调用

    """
    parser = argparse.ArgumentParser(description='score a model on a dataset')

    #model代表存放的模型的位置
    parser.add_argument('--model', type=str, default='model/resnet-18')

    #设置gpus,默认为0
    parser.add_argument('--gpus', type=str, default='0')

    #batch-size的次数，默认是10次
    parser.add_argument('--batch-size', type=int, default=10)

    #开始批次begin-epoch
    parser.add_argument('--begin-epoch', type=int, default=0)

    #图片的大小
    parser.add_argument('--image-shape', type=str, default='3,224,224')

    #修改训练集和验证集图像的尺寸
    parser.add_argument('--resize-train', type=int, default=256)
    parser.add_argument('--resize-val', type=int, default=224)

    #生成指定的训练数据，在生成.rec文件之间，必须要获取.lst文件，然后利用im2ec.py进行生成.rec文件
    parser.add_argument('--data-train-rec', type=str, default='data/data_train.rec')
    parser.add_argument('--data-train-idx', type=str, default='data/data_train.idx')
    parser.add_argument('--data-val-rec', type=str, default='data/data_val.rec')
    parser.add_argument('--data-val-idx', type=str, default='data/data_val.idx')

    #数据集的类别数，在今后的分类数据集中，这里的类数真的很重要，要时刻谨记
    parser.add_argument('--num-classes', type=int, default=2)

    #设置学习率；同样也可以自己指定想要的学习率
    parser.add_argument('--lr', type=float, default=0.005)

    #训练的epoch数，这里如果自己电脑配置很好，可以设置多一点，但是多了好像也不会有好的结果
    parser.add_argument('--num-epoch', type=int, default=4)
    parser.add_argument('--period', type=int, default=100)

    #模型保存的路径，根据实际情况设置
    parser.add_argument('--save-result', type=str, default='output/resnet-18/')

    #例子数
    parser.add_argument('--num-examples', type=int, default=22500)

    #动量和权重衰减参数的设置
    parser.add_argument('--mom', type=float, default=0.9, help='momentum for sgd')
    parser.add_argument('--wd', type=float, default=0.0001, help='weight decay for sgd')

    #保存的模型名字
    parser.add_argument('--save-name', type=str, default='resnet-18')

    #随机镜像，默认是1
    parser.add_argument('--random-mirror', type=int, default=0,
                        help='if or not randomly flip horizontally')

    #对比度的设置,默认0.3
    parser.add_argument('--max-random-contrast', type=float, default=0,
                        help='Chanege the contrast with a value randomly chosen from [-max, max]')

    #翻转，默认是15
    parser.add_argument('--max-rotate-angle', type=int, default=0,
                        help='Rotate by a random degree in [-v,v]')

    #网络层的名字全连接层前面的网络层名字
    parser.add_argument('--layer-name', type=str, default='flatten0',
                        help='the layer name before fullyconnected layer')

    #这里的是用于多少部乘以一个学习因子
    parser.add_argument('--factor', type=float, default=0.2, help='factor for learning rate decay')
    parser.add_argument('--step', type=int, default=5, help='step for learning rate decay')
    parser.add_argument('--from-scratch', type=bool, default=False,
                        help='Whether train from scratch')
    parser.add_argument('--fix-pretrain-param', type=bool, default=False,
                        help='Whether fix parameters of pretrain model')

    #这里通过parser对象的parse_args获取解析的参数
    args = parser.parse_args()
    return args

def main():
    #得到所有的配置参数
    args = parse_arguments()
    #判断是否存在 save_result
    if not os.path.exists(args.save_result):
        os.makedirs(args.save_result)

    #这里是设置日志信息的显示和保存
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    #显示管理对象
    stream_handler = logging.StreamHandler()
    logger.addHandler(stream_handler)

    #文件管理，用于保存信息的位置，我们可以自己指定
    file_handler = logging.FileHandler(args.save_result + '/train.log')
    logger.addHandler(file_handler)

    #将所有的参数信息保存到训练日志文件中
    logger.info(args)

    #训练模型的函数，然后将所有的参数都传进去，对应参数parse argument里的参数设置
    train_model(args=args)

if __name__ == '__main__':
    main()

测试：

import mxnet as mx
import numpy as np

"""
    1、调用load_model()函数导入预先训练好的模型
"""
def load_model(model_prefix, index, context, data_shapes, label_shapes):
    """
    1、通过load_checkpoint()函数进行读入模型：有两个参数：1、model存放的路径地址 2、训练模型是时得到的epoch信息index

    :param model_prefix:
    :param index:
    :param context:
    :param data_shapes:
    :param label_shapes:
    :return:
    """
    #通过此窗口获得到模型sym网络结构、arg_params表示网络的参数：对应是字典、aux_params是辅助参数，表示BN层的全局均值和方差
    sym, arg_params, aux_params = mx.model.load_checkpoint(model_prefix, index)

    #使用Module模块初始化一个module对象model
    model = mx.mod.Module(symbol=sym, context=context)

    """
        bind操作，将网络结构和输入数据绑定在一起
        输入数据的信息data_shape、输入数据的标签label_shape，两个都是一个元组组成的列表
        
    """
    model.bind(for_training=False,
               data_shapes=data_shapes,
               label_shapes=label_shapes)
    model.set_params(arg_params=arg_params,
                     aux_params=aux_params,
                     allow_missing=True)
    return model

"""
    2、load_data()函数读取的数据也应该做一定的预处理
"""
def load_data(data_path):
    #读取图像数据
    data = mx.image.imread(data_path)

    #将数值类型从int8转为float32
    cast_aug = mx.image.CastAug()

    #将图像resize到指定的尺寸
    resize_aug = mx.image.ForceResizeAug(size=[224, 224])

    #对输入图像的像素值做归一化处理;实现方法就是将原图像的像素值减去均值mean、然后除以标准差std得到返回值
    norm_aug = mx.image.ColorNormalizeAug(mx.nd.array([123, 117, 104]),
                                          mx.nd.array([1, 1, 1]))
    cla_augmenters = [cast_aug, resize_aug, norm_aug]

    for aug in cla_augmenters:
        data = aug(data)
    data = mx.nd.transpose(data, axes=(2, 0, 1))
    data = mx.nd.expand_dims(data, axis=0)
    data = mx.io.DataBatch([data])

    """
        输出的数据：将原始的[H,W,C]调整为[C,H,W]
        original shape：[224,224,3]
        DataBatch: data shapes: [(1, 3, 224, 224)] label shapes: None
    """
    return data

"""
    3、获取预测的结果;得到预测的类别概率，最后取预测概率最大的类别作为预测结果
"""
def get_output(model, data):
    model.forward(data)
    cla_prob = model.get_outputs()[0][0].asnumpy()
    cla_label = np.argmax(cla_prob)
    return cla_label

def main():
    #总共就两个类用0，1表示
    label_map = {0: "cat", 1: "dog"}

    #模型存放的地址
    model_prefix = "output/resnet-18/resnet-18"
    index = 4
    context = mx.gpu(0)
    data_shapes = [('data', (1, 3, 224, 224))]
    label_shapes = [('softmax_label', (1,))]

    #加载模型
    model = load_model(model_prefix, index, context, data_shapes, label_shapes)

    #测试图片的地址
    data_path = "data/demo_img2.jpg"
    data = load_data(data_path)

    cla_label = get_output(model, data)
    print("Predict result: {}".format(label_map.get(cla_label)))

if __name__ == '__main__':
    main()