行人重识别(12)——代码实践之模型训练（train_model(表征学习).py）

！转载请注明原文地址！——东方旅行者

更多行人重识别文章移步我的专栏：行人重识别专栏

表征学习模型训练（train_model(表征学习).py）

一、train_model(表征学习).py作用

本文件是行人重识别系统的核心文件，用于模型训练。

二、train_model(表征学习).py编写思路

为了便于管理一些常用参数，将其形成变量单独指定，这些参数有：

1.	输入图片宽度width
2.	输入图片高度height
3.	训练批次大小train_batch_size
4.	测试批次大小test_batch_size
5.	学习率train_lr
6.	开始训练的批次start_epoch
7.	结束训练的批次end_epoch
8.	动态学习率变化步长dy_step_size
9.	动态学习率变化倍数dy_step_gamma
10.	是否测试evaluate
11.	模型保存的地址best_model_path、final_model_path
12.	最大准确率（用于对判断最优模型）max_acc

1.train()方法

训练函数train()需要的参数有epoch（当前批次）、model（使用的模型）、criterion_class（损失函数）、optimizer（优化器类型，用于反向传播优化网络参数）、scheduler（用于管理学习率）、data_loader（指定数据加载器，获得网络的输入数据）。
首先使用enumerate迭代数据加载器吐出数据，每一次突出的数据是当前批数、（图片，行人ID，摄像机ID），在获取输入数据后，对优化器进行清零，防止上次计算结果对这次计算产生影响，使用模型进行运算，得到运算结果（表征学习阶段计算结果就是分类向量），使用损失函数根据计算结果与真实结果计算损失，然后使用损失的backward()方法对损失进行反向传播优化计算，然后使用scheduler的step()方法更新学习率，在损失的backward()计算完毕后，一定不要忘记使用optimizer的step()方法更新参数，否则网络参数不会修改。
然后使用torch.argmax按行求最大值，即计算分类结果，然后根据分类结果与真实结果计算本次训练的准确率。若本次准确率高于最高准确率，则保存此时模型至最优模型地址并打印相关信息。

2.main()方法

主函数main()是训练过程调用的函数，首先需要使用数据管理器加载数据集，便于之后数据加载器获取索引列表。然后进行训练数据处理器、测试数据处理器的声明。
训练数据处理器需要使用自定义的随机裁剪、水平翻转、将图片转为张量、归一化。
测试数据处理器需要使用图片尺度调整（统一尺寸）、将图片转为张量、归一化。
然后声明训练集数据吞吐器、测试集数据吞吐器与查询集数据吞吐器。在生成吞吐器需要指定使用的Dataset类型的数据集（自定义的ImageDataset，在dataset_loader.py中，同时在这里还需要指定使用的数据处理器），一个批次的大小，是否重排序，是否丢弃最后无法组成一整个批次的数据。
接下来加载模型，设置损失函数、优化器（torch.optim.SGD需要指定待优化的参数，学习率，权重衰减）、动态学习率（lr_scheduler.ReduceLROnPlateau需要指定使用的优化器；mode可选择‘min’或者‘max’，min表示当监控量停止下降的时候，学习率将减小，max表示当监控量停止上升的时候，学习率将减小。默认值为‘min’；factor学习率每次降低多少；patience容忍网路的性能不提升的次数，高于这个次数就降低学习率；min_lr学习率的下限）等。
设置模型为训练模式，调用train函数进行模型训练。最后将模型参数保存到目标地址。

三、代码

import os,sys,time,datetime
import os.path as osp
import numpy as np
from IPython import embed

import torch
import torch.nn as nn
import transform as T
from torch.utils.data import DataLoader
from torch.optim import lr_scheduler
from model.ReIDNet import ReIDNet
from dataset_manager import Market1501#数据管理器
from dataset_loader import ImageDataset#数据加载器

"""
本文件是行人重识别系统的核心文件，用于表征学习模型训练。
"""
#设定输入参数
width=64			#图片宽度
height=128			#图片高度
train_batch_size=32	#训练批量
test_batch_size=32	#测试批量
train_lr=0.01		#学习率
start_epoch=0		#开始训练的批次
end_epoch=1			#结束训练的批次
dy_step_size=800	#动态学习率变化步长
dy_step_gamma=0.9	#动态学习率变化倍数
evaluate=False		#是否测试
max_acc=-1			#最大准确率
best_model_path='./model/param/net_params_best.pth'		#最优模型保存地址
final_model_path='./model/param/net_params_final.pth'	#最终模型保存地址

def main():
    #数据集加载
    dataset=Market1501()
    
    #训练数据处理器
    transform_train=T.Compose([
        T.Random2DTransform(height,width),#尺度统一，随机裁剪
        T.RandomHorizontalFlip(),#水平翻转
        T.ToTensor(),#图片转张量
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),#归一化，参数固定
    ]
    )
    
    #测试数据处理器
    transform_test=T.Compose([
        T.Resize((height,width)),#尺度统一
        T.ToTensor(),#图片转张量
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),#归一化，参数固定
    ]
    )
    
    #train数据集吞吐器
    train_data_loader=DataLoader(
        ImageDataset(dataset.train, transform=transform_train),#自定义的数据集，使用训练数据处理器
        batch_size=train_batch_size,#一个批次的大小（一个批次有多少个图片张量）
        drop_last=True,#丢弃最后无法称为一整个批次的数据
    )
    print("train_data_loader inited")
    
    #query数据集吞吐器
    query_data_loader=DataLoader(
        ImageDataset(dataset.query, transform=transform_test),#自定义的数据集，使用测试数据处理器
        batch_size=test_batch_size,#一个批次的大小（一个批次有多少个图片张量）
        shuffle=False,#不重排
        drop_last=True,#丢弃最后无法称为一整个批次的数据
    )
    print("query_data_loader inited")
    
    #gallery数据集吞吐器
    gallery_data_loader=DataLoader(
        ImageDataset(dataset.gallery, transform=transform_test),#自定义的数据集，使用测试数据处理器
        batch_size=test_batch_size,#一个批次的大小（一个批次有多少个图片张量）
        shuffle=False,#不重排
        drop_last=True,#丢弃最后无法称为一整个批次的数据
    )
    print("gallery_data_loader inited\n")
    
    #加载模型
    model=ReIDNet(num_classes=751,loss={'softmax'})#指定分类的数量，与使用的损失函数以便决定模型输出何种计算结果
    print("=>ReIDNet loaded")
    print("Model size: {:.5f}M\n".format(sum(p.numel() for p in model.parameters())/1000000.0))
    
    #损失函数
    criterion_class=nn.CrossEntropyLoss()
    
    """
    优化器
    参数1，待优化的参数
    参数2，学习率
    参数3，权重衰减
    """
    optimizer=torch.optim.SGD(model.parameters(),lr=train_lr,weight_decay=5e-04)
    
    """
    动态学习率
    参数1，指定使用的优化器
    参数2，mode，可选择‘min’（min表示当监控量停止下降的时候，学习率将减小）或者‘max’（max表示当监控量停止上升的时候，学习率将减小）
    参数3，factor，代表学习率每次降低多少
    参数4，patience，容忍网路的性能不提升的次数，高于这个次数就降低学习率
    参数5，min_lr，学习率的下限
    """
    scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=dy_step_gamma, patience=10, min_lr=0.0001)
    
    #如果是测试
    if evaluate:
        test(model,query_data_loader,gallery_data_loader)
        return 0
    #如果是训练
    print('————model start training————\n')
    bt=time.time()#训练的开始时间
    for epoch in range(start_epoch,end_epoch):
        model.train(True)
        train(epoch,model,criterion_class,optimizer,scheduler,train_data_loader)
    et=time.time()#训练的结束时间
    print('**模型训练结束, 保存最终参数到{}**\n'.format(final_model_path))
    torch.save(model.state_dict(), final_model_path)
    print('————训练总用时{:.2f}小时————'.format((et-bt)/3600.0))

def train(epoch, model, criterion_class, optimizer, scheduler, data_loader):
    """
    训练函数train
    参数1，epoch（当前批次）
    参数2，model（使用的模型）
    参数3，criterion_class（损失函数）
    参数4，optimizer（优化器类型，用于反向传播优化网络参数）
    参数5，scheduler（用于管理学习率）
    参数6，data_loader（指定数据加载器，获得网络的输入数据）
    """
    global max_acc
    for batch_idx, (imgs, pids, cids) in enumerate(data_loader):
        optimizer.zero_grad()#优化器进行清零，防止上次计算结果对这次计算产生影响
        outputs=model(imgs)
        loss=criterion_class(outputs,pids)#使用损失函数根据计算结果与真实结果计算损失
        loss.backward()#根据损失进行反向传播优化计算
        scheduler.step(loss)#更新学习率
        optimizer.step()#更新网络中指定的需要优化的参数
        pred = torch.argmax(outputs, 1)#按行求最大值，计算分类结果
        current_acc=100*(pred == pids).sum().float()/len(pids)
        if current_acc>max_acc:
            max_acc=current_acc
            print('**最高准确度更新为{}%，保存此模型到{}**\n'.format(max_acc,best_model_path))
            torch.save(model.state_dict(), best_model_path)
        if batch_idx%10==0:
            print('————————————————————————————————')
            pred = torch.argmax(outputs, 1)
            print('Epoch: {}, Batch: {}, Loss: {}'.format(epoch + 1, batch_idx, loss.data))
            print('Current Accuracy: {:.2f}%'.format(100*(pred == pids).sum().float()/len(pids)))     
            print('————————————————————————————————\n')
    return 0

main()

四、训练打印日志

=> Market1501 loaded
------------------------------------------------------------------------
  subset: train  	| num_id:   751  	|  num_imgs:   12936  
  subset: query  	| num_id:   750  	|  num_imgs:    3368  
  subset: gallery 	| num_id:   751  	|  num_imgs:   19732  
------------------------------------------------------------------------
  total 			| num_id:  1501  	|  num_imgs:   16304  
------------------------------------------------------------------------
train_data_loader inited
query_data_loader inited
gallery_data_loader inited

=>ReIDNet loaded
Model size: 4.53483M

————model start training————

**最高准确度更新为0.0%，保存此模型到./model/param/net_params_best.pth**

————————————————————————————————
Epoch: 1, Batch: 0, Loss: 6.598439693450928
Current Accuracy: 0.00%
————————————————————————————————

**最高准确度更新为21.875%，保存此模型到./model/param/net_params_best.pth**

————————————————————————————————
Epoch: 1, Batch: 100, Loss: 6.623084545135498
Current Accuracy: 0.00%
————————————————————————————————
******************************中间省略若干******************************
******************************中间省略若干******************************
————————————————————————————————
————————————————————————————————
epoch: 50, batch: 400, loss: 5.373966217041016, lr:0.0009847709021836117
saving model params
model params saved
————————————————————————————————

saving model params
model params saved

————训练总用时5.08小时————