【CenterFusion】训练执行过程CenterFusion/src/main.py
文件作用：CenterFusion 项目训练的执行过程
- 首先查看文件的最末尾，最开始执行这段代码
if __name__ == '__main__':
opt = opts().parse()
'''
调用 opts 类中的 parse() 函数
这个函数在 CenterFusion/src/lib/opts.py 第 305 行
'''
main(opt)
'''
执行 main() 函数
'''
- 然后执行 main() 函数中的代码
def main(opt):
torch.manual_seed(opt.seed)
'''
这里引用了 torch 库内的函数 manual_seed() ：设置 CPU 生成随机数种子
seed 在 opts.py 文件第 50 行定义，默认值为 317
'''
torch.backends.cudnn.benchmark = not opt.not_cuda_benchmark and not opt.eval
'''
参数 not_cuda_benchmark 在 opts.py 第 48 行，含义：不是cuda基准
train.sh 脚本中没有添加 --not_cuda_benchmark 参数，所以 opt.not_cuda_benchmark 值为 False
参数 eval 在 opts.py 第 22 行，含义：只评估测试集 mini_val 并退出
train.sh 脚本中没有添加 --eval 参数，所以 opt.eval 值为 False
综上得出，torch.backends.cudnn.benchmark = True
PS:
设置为 True 会让程序在开始时花费一点额外时间，使用 cudnn 为整个网络的每个卷积层搜索最适合它的卷积实现算法，进而实现网络的加速
torch.backends.cudnn.benchmark是用于控制在使用CUDNN进行卷积操作时的优化级别的flag，该flag可以加速训练。如果该flag设置为True，则每次卷积操作都会针对输入大小进行优化，以便找到最快的卷积算法。但是，由于该操作会使用一些额外的内存，因此在使用较小的输入数据集时，并不总是最优的选择。
因此，如果你在使用较大的数据集进行训练，可以使用以下代码来启用优化级别：
torch.backends.cudnn.benchmark = True
但是，如果你在使用较小的数据集，并且想减少内存占用，可以使用以下代码来禁用优化级别：
torch.backends.cudnn.benchmark = False
'''
Dataset = get_dataset(opt.dataset)
'''
参数 dataset 在 opts.py 第 16 行，含义：设置默认数据集 nuscenes
get_dataset() 函数在 CenterFusion/src/lib/dataset/dataset_factory.py 第 32 行
返回一个 nuScenes 对象
这个 nuScenes 对象定义在 CenterFusion/src/lib/dataset/datasets/nuscenes.py 第 24 行
'''
opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
'''
更新一些配置信息
update_dataset_info_and_set_heads() 函数在 opts.py 第 458 行
'''
if not opt.not_set_cuda_env:
os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str
'''
opt.not_set_cuda_env 在脚本中没有添加这个参数，所以值为 False，再用 not 取反，则为 True
gpus_str 是 opt 中的一个 GPU 索引号字符串，如：'0,1'
这里是为了给系统添加 cuda 索引号
'''
opt.device = torch.device('cuda' if opt.gpus[0] >= 0 else 'cpu')
'''
设置模型要分配的位置
'''
logger = Logger(opt)
'''
新建一个 Logger 对象，保存 opt 的配置信息
Logger 对象在 CenterFusion/src/lib/logger.py
'''
print('Creating model...')
model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
'''
创建 DLA 模型（CenterNet 中的一种）
参数：
arch ：网络结构名称
heads ：网络的头部
head_conv ：头部输出的通道个数
create_model 函数在 CenterFusion/src/lib/model/model.py 第 24 行
'''
optimizer = get_optimizer(opt, model)
'''
根据超参数和模型定义优化器
'''
start_epoch = 0
'''
设定初始轮次
'''
lr = opt.lr
'''
获取设定的学习率：0.00025
'''
if opt.load_model != '':
model, optimizer, start_epoch = load_model(
model, opt.load_model, opt, optimizer)
'''
如果加载的预训练模型不为空，则加载模型
返回加载后的模型、优化器、初始轮次
load_model() 函数在 CenterFusion/src/lib/model/model.py 第 31 行
'''
trainer = Trainer(opt, model, optimizer)
'''
根据参数、模型、优化器得到对应任务的训练器（其中包括损失统计、模型损失等）
Trainer 类在 CenterFusion/src/lib/trainer.py 第 129 行
'''
trainer.set_device(opt.gpus, opt.chunk_sizes, opt.device)
'''
将训练器加载到 GPU 上
路径 CenterFusion/src/lib/trainer.py 第 137 行
'''
if opt.val_intervals < opt.num_epochs or opt.eval:
'''
val_intervals=1 < num_epochs=60 执行该 if 语句
'''
print('Setting up validation data...')
val_loader = torch.utils.data.DataLoader(
Dataset(opt, opt.val_split), batch_size=1, shuffle=False,
num_workers=1, pin_memory=True)
'''
加载测试集数据
torch.utils.data.DataLoader 是一个数据读取的一个接口，参数：
dataset (Dataset)：加载数据的数据集
batch_size (int, optional)：每个 batch 加载多少个样本(默认: 1)
shuffle (bool, optional)：设置为 True 时会在每个 epoch 重新打乱数据(默认: False)
num_workers (int, optional)：用多少个子进程加载数据。0 表示数据将在主进程中加载(默认: 0)
pin_memory (bool, optional)：设置 pin_memory=True，则意味着生成的 Tensor 数据最开始是属于内存中的锁页内存，
这样将内存的 Tensor 转义到 GPU 的显存就会更快一些
'''
if opt.eval:
'''
如果是测试，则执行，训练时，该 if 语句没有执行
'''
_, preds = trainer.val(0, val_loader)
'''
路径 CenterFusion/src/lib/trainer.py 第 402 行
'''
val_loader.dataset.run_eval(preds, opt.save_dir, n_plots=opt.eval_n_plots,
render_curves=opt.eval_render_curves)
'''
进行结果评判
'''
return
print('Setting up train data...')
train_loader = torch.utils.data.DataLoader(
Dataset(opt, opt.train_split), batch_size=opt.batch_size,
shuffle=opt.shuffle_train, num_workers=opt.num_workers,
pin_memory=True, drop_last=True
)
'''
加载训练集
drop_last (bool, optional) – 如果数据集大小不能被 batch size 整除，则设置为 True 后可删除最后一个不完整的 batch。
如果设为 False 并且数据集的大小不能被 batch size 整除，则最后一个 batch 将更小。(默认: False)
'''
print('Starting training...')
for epoch in range(start_epoch + 1, opt.num_epochs + 1):
'''
循环 60 次，训练 60 轮
'''
mark = epoch if opt.save_all else 'last'
'''
是否每 5 个 epoch 保存模型到磁盘
由于 train.sh 中没有添加 save_all 参数，所以为 False
最后 mark = 'last'，意为最后再保存模型到磁盘中
'''
for param_group in optimizer.param_groups:
lr = param_group['lr']
logger.scalar_summary('LR', lr, epoch)
break
'''
记录学习率
scalar_summary() 函数在 CenterFusion/src/lib/logger.py 第 73 行
'''
log_dict_train, _ = trainer.train(epoch, train_loader)
logger.write('epoch: {} |'.format(epoch))
'''
训练一轮模型，返回 ret 和 result
这里 log_dict_train = ret
train() 函数在 CenterFusion/src/lib/trainer.py 第 405 行
'''
for k, v in log_dict_train.items():
logger.scalar_summary('train_{}'.format(k), v, epoch)
logger.write('{} {:8f} | '.format(k, v))
'''
items() 函数：将一个字典以列表的形式返回，因为字典是无序的，所以返回的列表也是无序的
记录训练一次返回的结果
'''
if opt.val_intervals > 0 and epoch % opt.val_intervals == 0:
'''
val_intervals 的值为 1
epoch % opt.val_intervals 的值始终为 0
所以每训练一个 epoch 都要使用测试集进行一次测试
'''
save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(mark)),
epoch, model, optimizer)
'''
保存模型，下一轮训练好后又会覆盖上一轮的模型
save_model() 函数在 CenterFusion/src/lib/model/model.py 第 117 行
其中参数分别为保存路径、训练轮次、训练模型、优化器
最后保存的模型放在 ~/CenterFusion/src/lib/../../exp/ddd/centerfusion 文件夹下，后缀名为 .pth
'''
with torch.no_grad():
'''
with 语句适用于对资源进行访问的场合，确保不管使用过程中是否发生异常都会执行必要的“清理”操作，释放资源
比如文件使用后自动关闭／线程中锁的自动获取和释放等
with torch.no_grad(): 强制后面不进行计算图（计算过程的构建，以便梯度反向传播等操作）的构建
'''
log_dict_val, preds = trainer.val(epoch, val_loader)
'''
对当前轮次训练后的模型进行测试
'''
if opt.run_dataset_eval:
'''
run_dataset_eval 在 train.sh 中添加了该参数，所以值为 True
执行该 if 语句
'''
out_dir = val_loader.dataset.run_eval(preds, opt.save_dir,
n_plots=opt.eval_n_plots,
render_curves=opt.eval_render_curves)
'''
对测试结果评估
'''
with open('{}/metrics_summary.json'.format(out_dir), 'r') as f:
metrics = json.load(f)
logger.scalar_summary('AP/overall', metrics['mean_ap']*100.0, epoch)
for k,v in metrics['mean_dist_aps'].items():
logger.scalar_summary('AP/{}'.format(k), v*100.0, epoch)
for k,v in metrics['tp_errors'].items():
logger.scalar_summary('Scores/{}'.format(k), v, epoch)
logger.scalar_summary('Scores/NDS', metrics['nd_score'], epoch)
'''
记录测试数据集的评估指标
'''
for k, v in log_dict_val.items():
logger.scalar_summary('val_{}'.format(k), v, epoch)
logger.write('{} {:8f} | '.format(k, v))
'''
记录测试结果
'''
#保存这个检查点
else:
save_model(os.path.join(opt.save_dir, 'model_last.pth'),
epoch, model, optimizer)
logger.write('\n')
if epoch in opt.save_point:
save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(epoch)),
epoch, model, optimizer)
'''
save_point 在 train.sh 中的值为 20,40,50
所以当 epoch 为 20,40,50 时保存模型
'''
if epoch in opt.lr_step:
'''
lr_step 在 train.sh 中的值为 50
那么 epoch = 50 轮次时，更新学习率 lr
'''
lr = opt.lr * (0.1 ** (opt.lr_step.index(epoch) + 1))
print('Drop LR to', lr)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
logger.close()
- 其中有get_optimizer()函数，具体如下：
def get_optimizer(opt, model):
if opt.optim == 'adam':
optimizer = torch.optim.Adam(model.parameters(), opt.lr)
'''
执行该 if 语句，其中 torch.optim.Adam 是实现 Adam 算法的函数
'''
elif opt.optim == 'sgd':
print('Using SGD')
optimizer = torch.optim.SGD(
model.parameters(), opt.lr, momentum=0.9, weight_decay=0.0001)
else:
assert 0, opt.optim
return optimizer
'''
返回优化器
'''