SlowFast 训练相关源码解析

0. 前言

• 目标：
  • 模型输出的形式。
  • 数据集标签是的形式。
  • 损失函数在哪里计算，有哪些选择。
  • 训练相关细节：
    • hooks
    • loggings
    • tensorboard
    • lr
• 训练一般通过 tools/run_net.py 实现，其中调用了 tools/train_net.py 中的 def train(cfg) 方法。

1. 训练流程

• 第一步：初始化若干参数，包括日志参数、分布式训练参数、random seed、multigrid等。
• 第二步：构建模型，并统计模型中的参数数量以及计算量。
• 第三步：构建训练相关功能模块，包括优化器、预训练模型、tensorboard、性能指标监控等。
• 第四步：构建训练集/验证集。
• 第五步：实际训练。这一部分包括很多内容，稍微罗列下。
  • 设置 multigrid 参数。
  • 训练一个epoch，即遍历训练集一次计算损失函数、梯度下降、更新参数、更新tensorboard参数。
  • 更BN参数(Compute precise BN stats)
  • 保存模型。
  • 在验证集上进行验证。

2. 损失函数相关

• 关注的细节：
  • 输入数据形式
  • 模型输出形式
  • 损失函数选择与计算

2.1. 输入数据的形式：

• data_loader 的结果包括四个，frames, label, index, extra_data
• frames 的形式是 C, T, H, W，经过loader后应该就是 B, C, T, H, W.
• label分两种情况，对于ava形式是 [num_boxes, num_classes]，经过loader后就是 batch_size, num_boxes, num_classes，kinetics的形式是整数，经过loader后是[batch_size]。
• index是整数，经过loader后是 [batch_size]
• extrac_data 只有 ava有，kinetcis没有，是一个字典，包括 boxes, ori_boxes, metadata 三个key。

2.2. 模型输出形式：

• 模型分为两种，行为识别模型以及时空行为检测模型。
• 行为识别模型对应 slowfast/models/head_helper.py 中的 ResNetBasicHead。
  • 在该header中会先通过一个avg pooling，将输入的 N, C, T, H, W 转换为 N, C, 1, 1, 1 的形式，然后再transpose进行fc操作，最终得到[N, num_classes] 的logits。
  • 如果有多路分支会在ROI Pooling后将特征进行concat。
• 时空行为检测模型对应 slowfast/models/head_helper.py 中的 ResNetRoIHead。
  • 这部分操作先会根据T纬度进行avg pool，将输入从N, C, T, H, W转换为N, C, H, W，之后就是普通的ROI Pooling+fc判断每个ROI的类别，最终结果是 [num_boxes, num_classes]。
  • 如果有多路分支会在ROI Pooling后将特征进行concat。
  • boxes 是要外部输入（而不是算法内部生成）。

2.3. 损失函数选择与计算

• 通过 slowfast/models/losses.py 中的 get_loss_func 实现，包括了
  • nn.CrossEntropyLoss
  • nn.BCELoss
  • nn.BCEWithLogitsLoss

3. 性能指标/logging/tensorboard 相关

3.1. 性能指标(包括logging)

• AVA 使用 slowfast/utils/meters.py 中的 AVAMeter 记录相关性能指标。
• 行为识别数据集使用 slowfast/utils/meters.py 中的 TrainMeter 记录相关性能指标。
• 相关类的常用方法是：
  • iter_tic与iter_toc来设置开始、结束结束。
  • update_stats 更新记录的参数。
  • log_iter_stats 输出每iter的数据，有logging和本地json文件两种输出方式。
  • log_epoch_stats 输出每epoch的数据，有logging和本地json文件两种输出方式。
• log记录的内容包括：
  • _type：表示的是log数据类型，比如train_iter/val_iter/test_iter/val_epoch/test_epoch
  • cur_epoch
  • cur_iter
  • eta：训练剩余时间。计算方法是：当前 iter 时间 * (总iter数量 - 当前iter编号)
    • 训练的时候，这个数值存在较大问题：如果进行100次训练，其中一个iter耗时很长，剩余耗时很短，那 eta 展示的时间就很有迷惑性。
  • time_diff：当前 iter 使用时间（不知道为啥叫这个名字）
  • mode：train/val/test 其中之一
  • loss
  • lr
  • top1_err/top5_err：仅用于行为识别数据集。
• 记录一些需要累加算平均的数值时，使用了 slowfast/utils/meters.py 中的 ScalarMeter
• 其他用到的一些配置参数
  • cfg.LOG_PERIOD：log的频率
  • cfg.DATA.MULTI_LABEL：判断是否需要输出 top1/top5_err

3.3. tensorboard

• 主要实现代码都在 slowfast/utils/tensorboard_vis.py 中。
  • 基本功能都封装在 TensorboardWriter 中。
• 常用参数：
  • cfg.TENSORBOARD.ENABLE：是否使用tensorboard相关功能。默认关闭
  • cfg.TENSORBOARD.LOG_DIR：保存路径，空是保存在 cfg.OUTPUT_DIR 中。
• 所做的工作
  • 训练的时候就是记录loss/lr/top1_err/top5_err。
  • 每个iter都记录。

4. 多线程/分布式训练相关

• 之后补充