Realization of ReID code based on representation learning (5)

from __future__ import absolute_import
import os
import sys
import time
import datetime
import argparse
import os.path as osp
import numpy as np

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torch.backends.cudnn as cudnn
from torch.optim import lr_scheduler

import models
from utils.losses import CrossEntropyLabelSmooth
from utils import data_manager
from utils import transforms as T
from utils.data_loader import ImageDataset
from utils.utils import Logger
from utils.utils import AverageMeter, Logger, save_checkpoint
from utils.eval_metrics import evaluate
from utils.optimizers import init_optim

from IPython import embed

parser = argparse.ArgumentParser(description='Train reid with cross entropy loss')
# 数据集配置
parser.add_argument('--root', type=str, default=r'D:\dataset', help='root path to data directory')
parser.add_argument('-d', '--dataset', type=str, default='market1501')
parser.add_argument('-j', '--workers', default=0, type=int,
                    help='number of data loading workers')
parser.add_argument('--height', type=int, default=256)
parser.add_argument('--width', type=int, default=128)

# 由于cuhk这个数据集比较特殊，因此需要一些特殊配置
# CUHK03-specific setting
parser.add_argument('--cuhk03-labeled', action='store_true',
                    help="whether to use labeled images, if false, detected images are used (default: False)")
parser.add_argument('--cuhk03-classic-split', action='store_true',
                    help="whether to use classic split by Li et al. CVPR'14 (default: False)")
parser.add_argument('--use-metric-cuhk03', action='store_true',
                    help="whether to use cuhk03-metric (default: False)")

# 优化的配置
parser.add_argument('--labelsmooth',action='store_true') # 防止overfit的小技巧
parser.add_argument('--optim',type=str, default='adam')
parser.add_argument('--max-epoch',default=20,type=int)
parser.add_argument('--start-epoch', default=0,type=int)
parser.add_argument('--train-batch', default=32, type=int)
parser.add_argument('--test-batch',default=32,type=int)
parser.add_argument('--lr', '--learning-rate', default=0.0002, type=float)
parser.add_argument('--stepsize', default=5,type=int) #多少个epoch后衰减一次学习率1
parser.add_argument('--gamma',default=0.1, type=float) #每次学习率衰减多少，0.1表示每次降低为原来的十分之一
parser.add_argument('--weight-decay', default=5e-04,type=float) #正则化参数，防止overfit

# 架构
parser.add_argument('-a','--arch', type=str, default='resnet50')

# 训练过程的参数
parser.add_argument('--print-freq', type=int, default=5) #训练多少个epoch打印一次
parser.add_argument('--seed', type=int, default=1) # 随机化参数，保证每次结果都能稳定地复现
parser.add_argument('--resume', type=str, default='', ) # 从哪个存档的模型开始恢复
parser.add_argument('--evaluate', action='store_true') #是否只用于测试
parser.add_argument('--eval-step', type=int, default=-1,
                    help="run evaluation for every N epochs (set to -1 to test after training)") #训练多少个epoch开始测试一次
parser.add_argument('--start-eval', type=int, default=0, help="start to evaluate after specific epoch") #从第几个epoch开始做测试
parser.add_argument('--save-dir', type=str, default='log')
parser.add_argument('--use-cpu', action='store_true', help="use cpu")
parser.add_argument('--gpu-devices', default='0', type=str, help='gpu device ids for CUDA_VISIBLE_DEVICES')

args = parser.parse_args()


def main():
    use_gpu = torch.cuda.is_available()
    if args.use_cpu:
        use_gpu = False
    if use_gpu:
        pin_memory = True # 在利用gpu时节省内存用的，
    else:
        pin_memory = False
    if not args.evaluate:
        sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
    else:
        sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
    print("==========\nArgs:{}\n==========".format(args))

    if use_gpu:
        print("Currently using GPU {}".format(args.gpu_devices))
        # 指定配置中选择使用的GPU
        os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
        cudnn.benchmark = True #用这句话就会把cudnn库用进来，使卷积算的更快
        torch.cuda.manual_seed_all(args.seed) # 把随机种子传进来，保证每次结果都能基本稳定复现
    else:
        print("Currently using CPU (GPU is highly recommended)")

    # 初始化dataset
    dataset = data_manager.init_img_dataset(root = args.root, name=args.dataset)# 后面两个参数应用于cuhk数据集

    # dataloader & augmentation train query gallery
    # 三个dataloader的区别在于训练时需要数据增广时不同
    transform_train = T.Compose([
        T.Random2DTranslation(args.height,args.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(args.height, args.width),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),

    ])
    trainloader = DataLoader(
        ImageDataset(dataset.train, transform=transform_train),
        batch_size = args.train_batch,num_workers=args.workers,
        pin_memory=pin_memory,drop_last=True,
    )
    queryloader = DataLoader(
        ImageDataset(dataset.query, transform=transform_test),
        batch_size=args.test_batch, num_workers=args.workers,
        shuffle=False,
        pin_memory=pin_memory, drop_last=True,
    )

    galleryloader = DataLoader(
        ImageDataset(dataset.gallery, transform=transform_test),
        batch_size=args.test_batch, num_workers=args.workers,
        shuffle=False,
        pin_memory=pin_memory, drop_last=True,
    )

    # 初始化模型
    print("Initializing model: {}".format(args.arch))
    model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'cent'})
    # 输出模型大小
    print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters()) / 1000000.0))

    # 定义损失
    #分类损失
    criterion_class = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu)

    # 优化器
    optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)# model.parameters()表示更新所有参数


    # 设计学习率衰减
    if args.stepsize > 0:
        # 阶梯型衰减
        scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize,gamma=args.gamma)

    start_epoch = args.start_epoch

    # 是否恢复模型
    if args.resume:
        print("Loading checkpoint from '{}'".format(args.resume))
        checkpoint = torch.load(args.resume)
        # 'state_dict' pytorch存模型，名字为层的名字，内容为层的参数
        model.load_state_dict(checkpoint['state_dict'])
        start_epoch = checkpoint['epoch']

    if use_gpu:
        # 使用并行库，多卡训练
        model = nn.DataParallel(model).cuda()

    # 若仅用于测试
    if args.evaluate:
        print("Evaluate only!")
        test(model, queryloader,galleryloader,use_gpu)
        return 0

    print("start training**************")

    for epoch in range(start_epoch, args.max_epoch):
        train(epoch, model, criterion_class, optimizer, trainloader, use_gpu)

def train(epoch, model, criterion_class, optimizer, trainloader, use_gpu):

    model.train() # 将模型设置为训练模式

    for batch_idx, (imgs, pids, _) in enumerate(trainloader):
        if use_gpu:
            imgs, pids = imgs.cuda(), pids.cuda()
        # 算前传
        outputs = model(imgs)
        loss = criterion_class(outputs, pids)
        # 清空梯度
        optimizer.zero_grad()
        # 反传播
        loss.backward()
        # 更新模型参数
        optimizer.step()
        print(batch_idx, loss)
        break


def test(model, queryloader, galleryloader, use_gpu, ranks=[1, 5, 10, 20]):
    batch_time = AverageMeter()

    model.eval()

    with torch.no_grad():
        qf, q_pids, q_camids = [], [], []
        for batch_idx, (imgs, pids, camids) in enumerate(queryloader):
            if use_gpu: imgs = imgs.cuda()

            end = time.time()
            features = model(imgs)
            batch_time.update(time.time() - end)

            features = features.data.cpu()
            qf.append(features)
            q_pids.extend(pids)
            q_camids.extend(camids)
        qf = torch.cat(qf, 0)
        q_pids = np.asarray(q_pids)
        q_camids = np.asarray(q_camids)

        print("Extracted features for query set, obtained {}-by-{} matrix".format(qf.size(0), qf.size(1)))

        gf, g_pids, g_camids = [], [], []
        for batch_idx, (imgs, pids, camids) in enumerate(galleryloader):
            if use_gpu: imgs = imgs.cuda()

            end = time.time()
            features = model(imgs)
            batch_time.update(time.time() - end)

            features = features.data.cpu()
            gf.append(features)
            g_pids.extend(pids)
            g_camids.extend(camids)
        gf = torch.cat(gf, 0)
        g_pids = np.asarray(g_pids)
        g_camids = np.asarray(g_camids)

        print("Extracted features for gallery set, obtained {}-by-{} matrix".format(gf.size(0), gf.size(1)))

    print("==> BatchTime(s)/BatchSize(img): {:.3f}/{}".format(batch_time.avg, args.test_batch))

    m, n = qf.size(0), gf.size(0)
    distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
              torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
    distmat.addmm_(1, -2, qf, gf.t())
    distmat = distmat.numpy()

    print("Computing CMC and mAP")
    cmc, mAP = evaluate(distmat, q_pids, g_pids, q_camids, g_camids, use_metric_cuhk03=args.use_metric_cuhk03)

    print("Results ----------")
    print("mAP: {:.1%}".format(mAP))
    print("CMC curve")
    for r in ranks:
        print("Rank-{:<3}: {:.1%}".format(r, cmc[r - 1]))
    print("------------------")

    return cmc[0]


if __name__ == '__main__':
    main()