TSN源码阅读

一、项目结构

1.py文件解释

main.py 训练脚本
test_models.py 测试脚本
opts.py 参数配置脚本
dataset.py 数据读取脚本
models.py 网络结构构建脚本
transforms.py 数据预处理相关脚本
tf_model_zoo 文件夹关于导入模型结构的脚本

2.函数组成及调用关系

3.IPO图

下图展示了，TSN如何将UCF-101数据集提出的帧进行分类的过程，标注了每一个Tensor的大小

二、opts.py解读

import argparse
parser = argparse.ArgumentParser(description="PyTorch implementation of Temporal Segment Networks")
parser.add_argument('dataset', type=str, choices=['ucf101', 'hmdb51', 'kinetics'])     三种类型的dataset
parser.add_argument('modality', type=str, choices=['RGB', 'Flow', 'RGBDiff'])         # 三种不同的输入
parser.add_argument('train_list', type=str)
parser.add_argument('val_list', type=str)

# ========================= Model Configs ==========================                   模型参数
parser.add_argument('--arch', type=str, default="resnet101")                         # 基本模型 默认resnet101
parser.add_argument('--num_segments', type=int, default=3)                           # segment的数值，默认3
parser.add_argument('--consensus_type', type=str, default='avg',                     # 聚合函数的选择，默认均值，choices avg、max、topk、identity、rnn、cnn
                    choices=['avg', 'max', 'topk', 'identity', 'rnn', 'cnn'])
parser.add_argument('--k', type=int, default=3)                                      # k的数值 默认3

parser.add_argument('--dropout', '--do', default=0.5, type=float,                    # dropout默认值 0.5
                    metavar='DO', help='dropout ratio (default: 0.5)')               # metavar大概可以理解为注释？
parser.add_argument('--loss_type', type=str, default="nll",
                    choices=['nll'])

# ========================= Learning Configs ==========================
parser.add_argument('--epochs', default=45, type=int, metavar='N',                   # 学习轮数默认45
                    help='number of total epochs to run')
parser.add_argument('-b', '--batch-size', default=256, type=int,                     # batch_size默认256
                    metavar='N', help='mini-batch size (default: 256)')
parser.add_argument('--lr', '--learning-rate', default=0.001, type=float,            # learning_rate默认0.001
                    metavar='LR', help='initial learning rate')
parser.add_argument('--lr_steps', default=[20, 40], type=float, nargs="+",           # 学习率每一轮下降10%
                    metavar='LRSteps', help='epochs to decay learning rate by 10')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',              # 冲量默认为0.9
                    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=5e-4, type=float,              # weight_dacay默认为5e-4
                    metavar='W', help='weight decay (default: 5e-4)')
parser.add_argument('--clip-gradient', '--gd', default=None, type=float,             # clip_gradient默认为none
                    metavar='W', help='gradient norm clipping (default: disabled)')
parser.add_argument('--no_partialbn', '--npb', default=False, action="store_true")   # no_partiblbn 应该就是论文中说的部分BN，初始值为false

# ========================= Monitor Configs ==========================
parser.add_argument('--print-freq', '-p', default=20, type=int,                      # 输出的频率 默认值20次一输出
                    metavar='N', help='print frequency (default: 20)')
parser.add_argument('--eval-freq', '-ef', default=5, type=int,                       #
                    metavar='N', help='evaluation frequency (default: 5)')


# ========================= Runtime Configs ==========================                # 用到再说吧，具体跑的时候设置的参数
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
                    help='evaluate model on validation set')
parser.add_argument('--snapshot_pref', type=str, default="")
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('--gpus', nargs='+', type=int, default=None)
parser.add_argument('--flow_prefix', default="", type=str)

三、main.py解读

1.整体架构

1.from opts import parser，解析命令行中的参数
2.调用models.py初始化TSN模型
3.调用dataset.py导入数据
4.训练、保存模型

2.具体代码解读

import argparse
import os
import time
import shutil
import torch
import torchvision
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
from torch.nn.utils import clip_grad_norm
from dataset import TSNDataSet
from models import TSN
from transforms import *
from opts import parser

"导入一些要用的包，其中比较重要的是 "
"导入模型：from models import TSN."
"导入配置的参数：from opts import parser."

# 最好的预测正确率！
best_prec1 = 0


def main():
    # global 全局变量
    global args, best_prec1
    # 进行cmd调参，具体设计的参数都在opts.py中
    args = parser.parse_args()

    if args.dataset == 'ucf101':
        num_class = 101
    elif args.dataset == 'hmdb51':
        num_class = 51
    elif args.dataset == 'kinetics':
        num_class = 400
    else:
        raise ValueError('Unknown dataset '+args.dataset)
    # 初始化模型
    model = TSN(num_class, args.num_segments, args.modality,base_model=args.arch,consensus_type=args.consensus_type, dropout=args.dropout, partial_bn=not args.no_partialbn)
    "TSN的定义在models.py脚本中"
    "num_class：分类的类别数"
    "args.num_segments：把一个video分成多少份，对应论文中的K，默认为K=3"
    "args.modality：采用哪种输入，比如RGB表示常规图像，Flow表示optical flow等"
    "args.arch：采用哪种模型，比如ResNet101，BNInception等"
    "rags.consensus_type：采用不同snippet融合方式，比如avg"
    "args.dropout:dropout参数"


    crop_size = model.crop_size
    scale_size = model.scale_size
    input_mean = model.input_mean
    input_std = model.input_std
    policies = model.get_optim_policies()
    train_augmentation = model.get_augmentation()

    # 多GPU与断点恢复设置，单机多卡
    # 使用torch.nn.DataParallel方法设置多GPU训练
    model = torch.nn.DataParallel(model, device_ids=args.gpus).cuda()

    # 而args.resume主要是用来设置是否从断点处继续训练，比如原来训练模型训到一半停止了，
    # 希望继续从保存的最新epoch开始训练，因此args.resume要么是默认的None，要么就是保存的模型文件（.pth）的路径
    if args.resume:
        if os.path.isfile(args.resume):
            print(("=> loading checkpoint '{}'".format(args.resume)))
            # 其中checkpoint = torch.load(args.resume)是用来导入已训练好的模型，
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            best_prec1 = checkpoint['best_prec1']
            # model.load_state_dict方法是完成导入模型的参数初始化model这个网络的过程，
            # 这也是torch.nn.Module类中的重要方法之一。
            model.load_state_dict(checkpoint['state_dict'])
            print(("=> loaded checkpoint '{}' (epoch {})"
                  .format(args.evaluate, checkpoint['epoch'])))
        else:
            print(("=> no checkpoint found at '{}'".format(args.resume)))

    cudnn.benchmark = True

    # 数据加载，通过自定义的TSNDataSet类导入数据
    if args.modality != 'RGBDiff':
        normalize = GroupNormalize(input_mean, input_std)
    else:
        normalize = IdentityTransform()

    if args.modality == 'RGB':
        data_length = 1
    elif args.modality in ['Flow', 'RGBDiff']:
        data_length = 5
    train_loader = torch.utils.data.DataLoader(
        TSNDataSet("", args.train_list, num_segments=args.num_segments,
                   new_length=data_length,
                   modality=args.modality,
                   image_tmpl="img_{:05d}.jpg" if args.modality in ["RGB", "RGBDiff"] else args.flow_prefix+"{}_{:05d}.jpg",
                   transform=torchvision.transforms.Compose([
                       train_augmentation,
                       Stack(roll=args.arch == 'BNInception'),
                       ToTorchFormatTensor(div=args.arch != 'BNInception'),
                       normalize,
                   ])),
        batch_size=args.batch_size, shuffle=True,
        num_workers=args.workers, pin_memory=True)

    val_loader = torch.utils.data.DataLoader(
        TSNDataSet("", args.val_list, num_segments=args.num_segments,
                   new_length=data_length,
                   modality=args.modality,
                   image_tmpl="img_{:05d}.jpg" if args.modality in ["RGB", "RGBDiff"] else args.flow_prefix+"{}_{:05d}.jpg",
                   random_shift=False,
                   transform=torchvision.transforms.Compose([
                       GroupScale(int(scale_size)),
                       GroupCenterCrop(crop_size),
                       Stack(roll=args.arch == 'BNInception'),
                       ToTorchFormatTensor(div=args.arch != 'BNInception'),
                       normalize,
                   ])),
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    # define loss function (criterion) and optimizer
    # 定义损失函数，优化器和设置一些超参数，
    # 从代码中可以看到，这里使用的是交叉熵损失函数，优化器使用SGD方式
    if args.loss_type == 'nll':
        criterion = torch.nn.CrossEntropyLoss().cuda()
    else:
        raise ValueError("Unknown loss type")

    for group in policies:
        print(('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format(
            group['name'], len(group['params']), group['lr_mult'], group['decay_mult'])))

    optimizer = torch.optim.SGD(policies,
                                args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    # 根据args.evaluate参数判断当前是训练模式还是测试模式
    if args.evaluate:
        # 是测试模型，直接调用validate方法验证模型
        validate(val_loader, model, criterion, 0)
        return
    # 不是测试模型，是训练模型
    for epoch in range(args.start_epoch, args.epochs):
        # 调用方法，调整学习率
        adjust_learning_rate(optimizer, epoch, args.lr_steps)

        # train for one epoch
        # 调用方法，开始训练模型
        train(train_loader, model, criterion, optimizer, epoch)

        # evaluate on validation set 模型验证和保存
        # 当训练epoch到达指定值，进行模型验证保存，使用args.eval_freq控制保存的epoch值
        if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
            prec1 = validate(val_loader, model, criterion, (epoch + 1) * len(train_loader))

            # remember best prec@1 and save checkpoint
            # 比较得到val更好的模型并保存
            is_best = prec1 > best_prec1
            best_prec1 = max(prec1, best_prec1)
            # 调用save_checkpoint保存模型参数和其他信息
            save_checkpoint({
   
                'epoch': epoch + 1,
                'arch': args.arch,
                'state_dict': model.state_dict(),
                'best_prec1': best_prec1,
            }, is_best)

# train函数是整个训练部分的入口，将TSN的训练部分封装起来，成为一个函数，之后在主函数直接调用即可
def train(train_loader, model, criterion, optimizer, epoch):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    # 自定义是否进行部分BN
    if args.no_partialbn:
        model.module.partialBN(False)
    else:
        model.module.partialBN(True)

    # 调用model.py中的train方法来对模型的参数进行预训练，
    # 并且冻结除第一层之外的所有批处理规范化层的均值和方差
    # 对全连接层的参数进行训练，达到微调的目的。
    model.train()

    # 运行数据迭代读取的循环函数
    # 当执行enumerate(train_loader)的时候，是先调用DataLoader类的__iter__方法，
    # 该方法里面再调用DataLoaderIter类的初始化操作__init__
    end = time.time()
    for i, (input, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        # 对train_loader中的数据集进行遍历，存储其中的数据集输入和真实标签
        target = target.cuda()
        input_var = torch.autograd.Variable(input)
        target_var = torch.autograd.Variable(target)

        # 执行output=model(input_var)得到模型的输入结果
        output = model(input_var)
        loss = criterion(output, target_var)

        # 调用计算损失函数之前调用accuracy函数来更新top1和top5的准确率
        prec1, prec5 = accuracy(output.data, target, topk=(1,5))
        losses.update(loss.data[0], input.size(0))
        top1.update(prec1[0], input.size(0))
        top5.update(prec5[0], input.size(0))


        # 对于梯度清零、回传和之前学习的过程一样，
        # 最后可以得到训练之后的模型和训练的loss、accuracy等
        optimizer.zero_grad()

        loss.backward()

        if args.clip_gradient is not None:
            total_norm = clip_grad_norm(model.parameters(), args.clip_gradient)
            if total_norm > args.clip_gradient:
                print("clipping gradient: {} with coef {}".format(total_norm, args.clip_gradient / total_norm))

        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:
            print(('Epoch: [{0}][{1}/{2}], lr: {lr:.5f}\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                   epoch, i, len(train_loader), batch_time=batch_time,
                   data_time=data_time, loss=losses, top1=top1, top5=top5, lr=optimizer.param_groups[-1]['lr'])))


# 验证函数validate和训练函数train类似
# 不同点是：model.eval()将模型设置为evaluate mode
# 没有optimizer.zero_grad()、loss.backward()、optimizer.step()等损失回传或梯度更新操作
def validate(val_loader, model, criterion, iter, logger=None):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    # switch to evaluate mode
    model.eval()

    end = time.time()
    for i, (input, target)