Question

在讲解主程序 main.py 之前，我们先来看看 2017 年 3 月谷歌开源的一个命令行工具 fire^[2] ，通过 pip install fire 即可安装。下面来看看 fire 的基础用法，假设 example.py 文件内容如下：

import fire

def add(x, y):
  return x + y
  
def mul(**kwargs):
    a = kwargs['a']
    b = kwargs['b']
    return a * b

if __name__ == '__main__':
  fire.Fire()

那么我们可以使用：

python example.py add 1 2 # 执行 add(1, 2)
python example.py mul --a=1 --b=2 # 执行 mul(a=1, b=2), kwargs={'a':1, 'b':2}
python example.py add --x=1 --y==2 # 执行 add(x=1, y=2)

可见，只要在程序中运行 fire.Fire() ，即可使用命令行参数 python file <function> [args,] {--kwargs,} 。fire 还支持更多的高级功能，具体请参考官方指南^[3] 。

在主程序 main.py 中，主要包含四个函数，其中三个需要命令行执行， main.py 的代码组织结构如下：

def train(**kwargs):
    '''
    训练
    '''
    pass
	 
def val(model, dataloader):
    '''
    计算模型在验证集上的准确率等信息，用以辅助训练
    '''
    pass

def test(**kwargs):
    '''
    测试（inference）
    '''
    pass

def help():
    '''
    打印帮助的信息 
    '''
    print('help')

if __name__=='__main__':
    import fire
    fire.Fire()

根据 fire 的使用方法，可通过 python main.py <function> --args=xx 的方式来执行训练或者测试。

训练

训练的主要步骤如下：

• 定义网络
• 定义数据
• 定义损失函数和优化器
• 计算重要指标
• 开始训练
  • 训练网络
  • 可视化各种指标
  • 计算在验证集上的指标

训练函数的代码如下：

def train(**kwargs):
    
    # 根据命令行参数更新配置
    opt.parse(kwargs)
    vis = Visualizer(opt.env)
    
    # step1: 模型
    model = getattr(models, opt.model)()
    if opt.load_model_path:
        model.load(opt.load_model_path)
    if opt.use_gpu: model.cuda()

    # step2: 数据
    train_data = DogCat(opt.train_data_root,train=True)
    val_data = DogCat(opt.train_data_root,train=False)
    train_dataloader = DataLoader(train_data,opt.batch_size,
                        shuffle=True,
                        num_workers=opt.num_workers)
    val_dataloader = DataLoader(val_data,opt.batch_size,
                        shuffle=False,
                        num_workers=opt.num_workers)
    
    # step3: 目标函数和优化器
    criterion = t.nn.CrossEntropyLoss()
    lr = opt.lr
    optimizer = t.optim.Adam(model.parameters(),
                            lr = lr,
                            weight_decay = opt.weight_decay)
        
    # step4: 统计指标：平滑处理之后的损失，还有混淆矩阵
    loss_meter = meter.AverageValueMeter()
    confusion_matrix = meter.ConfusionMeter(2)
    previous_loss = 1e100

    # 训练
    for epoch in range(opt.max_epoch):
        
        loss_meter.reset()
        confusion_matrix.reset()

        for ii,(data,label) in enumerate(train_dataloader):

            # 训练模型参数 
            input = Variable(data)
            target = Variable(label)
            if opt.use_gpu:
                input = input.cuda()
                target = target.cuda()
            optimizer.zero_grad()
            score = model(input)
            loss = criterion(score,target)
            loss.backward()
            optimizer.step()
            
            # 更新统计指标以及可视化
            loss_meter.add(loss.data[0])
            confusion_matrix.add(score.data, target.data)

            if ii%opt.print_freq==opt.print_freq-1:
                vis.plot('loss', loss_meter.value()[0])
                
                # 如果需要的话，进入 debug 模式
                if os.path.exists(opt.debug_file):
                    import ipdb;
                    ipdb.set_trace()

        model.save()

        # 计算验证集上的指标及可视化
        val_cm,val_accuracy = val(model,val_dataloader)
        vis.plot('val_accuracy',val_accuracy)
        vis.log("epoch:{epoch},lr:{lr},loss:{loss},train_cm:{train_cm},val_cm:{val_cm}"
        .format(
                    epoch = epoch,
                    loss = loss_meter.value()[0],
                    val_cm = str(val_cm.value()),
                    train_cm=str(confusion_matrix.value()),
                    lr=lr))
        
        # 如果损失不再下降，则降低学习率
        if loss_meter.value()[0] > previous_loss:          
            lr = lr * opt.lr_decay
            for param_group in optimizer.param_groups:
                param_group['lr'] = lr
                
        previous_loss = loss_meter.value()[0]

这里用到了 PyTorchNet^[4] 里面的一个工具: meter。meter 提供了一些轻量级的工具，用于帮助用户快速统计训练过程中的一些指标。 AverageValueMeter 能够计算所有数的平均值和标准差，这里用来统计一个 epoch 中损失的平均值。 confusionmeter 用来统计分类问题中的分类情况，是一个比准确率更详细的统计指标。例如对于表格 6-1，共有 50 张狗的图片，其中有 35 张被正确分类成了狗，还有 15 张被误判成猫；共有 100 张猫的图片，其中有 91 张被正确判为了猫，剩下 9 张被误判成狗。相比于准确率等统计信息，混淆矩阵更能体现分类的结果，尤其是在样本比例不均衡的情况下。

表 6-1 混淆矩阵

样本	判为狗	判为猫
实际是狗	35	15
实际是猫	9	91

PyTorchNet 从 TorchNet^[5] 迁移而来，提供了很多有用的工具，但其目前开发和文档都还不是很完善，本书不做过多的讲解。

验证

验证相对来说比较简单，但要注意需将模型置于验证模式( model.eval() )，验证完成后还需要将其置回为训练模式( model.train() )，这两句代码会影响 BatchNorm 和 Dropout 等层的运行模式。验证模型准确率的代码如下。

def val(model,dataloader):
    '''
    计算模型在验证集上的准确率等信息
    '''

    # 把模型设为验证模式
    model.eval()
    
    confusion_matrix = meter.ConfusionMeter(2)
    for ii, data in enumerate(dataloader):
        input, label = data
        val_input = Variable(input, volatile=True)
        val_label = Variable(label.long(), volatile=True)
        if opt.use_gpu:
            val_input = val_input.cuda()
            val_label = val_label.cuda()
        score = model(val_input)
        confusion_matrix.add(score.data.squeeze(), label.long())

    # 把模型恢复为训练模式
    model.train()
    
    cm_value = confusion_matrix.value()
    accuracy = 100. * (cm_value[0][0] + cm_value[1][1]) /\
			     (cm_value.sum())
    return confusion_matrix, accuracy

测试

测试时，需要计算每个样本属于狗的概率，并将结果保存成 csv 文件。测试的代码与验证比较相似，但需要自己加载模型和数据。

def test(**kwargs):
    opt.parse(kwargs)
    
    # 模型
    model = getattr(models, opt.model)().eval()
    if opt.load_model_path:
        model.load(opt.load_model_path)
    if opt.use_gpu: model.cuda()

    # 数据
    train_data = DogCat(opt.test_data_root,test=True)
    test_dataloader = DataLoader(train_data,\
							    batch_size=opt.batch_size,\
							    shuffle=False,\
							    num_workers=opt.num_workers)
    
    results = []
    for ii,(data,path) in enumerate(test_dataloader):
        input = t.autograd.Variable(data,volatile = True)
        if opt.use_gpu: input = input.cuda()
        score = model(input)
        probability = t.nn.functional.softmax\
	        (score)[:,1].data.tolist()      
        batch_results = [(path_,probability_) \
	        for path_,probability_ in zip(path,probability) ]
        results += batch_results
    write_csv(results,opt.result_file)
    return results

帮助函数

为了方便他人使用, 程序中还应当提供一个帮助函数，用于说明函数是如何使用。程序的命令行接口中有众多参数，如果手动用字符串表示不仅复杂，而且后期修改 config 文件时，还需要修改对应的帮助信息，十分不便。这里使用了 Python 标准库中的 inspect 方法，可以自动获取 config 的源代码。help 的代码如下:

def help():
    '''
    打印帮助的信息： python file.py help
    '''
    
    print('''
    usage : python {0} <function> [--args=value,]
    <function> := train | test | help
    example: 
            python {0} train --env='env0701' --lr=0.01
            python {0} test --dataset='path/to/dataset/root/'
            python {0} help
    avaiable args:'''.format(__file__))

    from inspect import getsource
    source = (getsource(opt.__class__))
    print(source)

当用户执行 python main.py help 的时候，会打印如下帮助信息：

    usage : python main.py <function> [--args=value,]
    <function> := train | test | help
    example: 
            python main.py train --env='env0701' --lr=0.01
            python main.py test --dataset='path/to/dataset/'
            python main.py help
    avaiable args:
class DefaultConfig(object):
    env = 'default' # visdom 环境
    model = 'AlexNet' # 使用的模型
    
    train_data_root = './data/train/' # 训练集存放路径
    test_data_root = './data/test1' # 测试集存放路径
    load_model_path = 'checkpoints/model.pth' # 加载预训练的模型

    batch_size = 128 # batch size
    use_gpu = True # user GPU or not
    num_workers = 4 # how many workers for loading data
    print_freq = 20 # print info every N batch

    debug_file = '/tmp/debug' 
    result_file = 'result.csv' # 结果文件
      
    max_epoch = 10
    lr = 0.1 # initial learning rate
    lr_decay = 0.95 # when val_loss increase, lr = lr*lr_decay
    weight_decay = 1e-4 # 损失函数