PyTorch Batch Norm使用训练mean/std及参数相关问题咨询

问题解答 #

1. 如何使用训练阶段的滑动均值/方差而非batch统计量 #

PyTorch BatchNorm层的计算逻辑由两个参数控制：模型当前是训练/评估模式，以及BatchNorm初始化时的track_running_stats参数（默认值为True）。
只需在需要使用训练累计统计量的阶段调用model.eval()，BatchNorm层就会自动使用训练阶段累计的running_mean和running_var做归一化，不会计算当前batch的统计量，可解决batch统计导致的模型发散问题。
如果是MAML这类嵌套训练的场景，若内层适配循环需要保持训练模式更新其他参数，但不想使用batch统计、也不想更新全局滑动统计量，可临时将所有BatchNorm层的momentum设为0，或是临时把track_running_stats设为False，避免内层操作污染全局的滑动统计结果。

2. 未训练时running_mean为全0属于正常初始化 #

PyTorch的BatchNorm层默认初始化时，running_mean的初始值就是全0向量，running_var的初始值为全1向量，这是框架的默认逻辑，并非参数异常，训练过程中会随输入batch不断更新这两个值。
如果训练完成后running_mean仍为全0，优先排查训练过程中是否正确调用了model.train()：只有模型处于训练模式时，才会正常更新滑动统计值，若训练全程模型都处于评估模式，滑动统计值会一直停留在初始值不会更新。

3. 滑动统计参数默认会保存在检查点文件中 #

running_mean、running_var属于模型的持久化缓冲区参数，默认调用torch.save(model.state_dict(), 检查点路径)时，会自动将这两个参数和其他可训练参数一起保存；加载时调用model.load_state_dict(torch.load(检查点路径))也会正常恢复这两个值，不需要额外操作。只有你手动在保存时过滤了缓冲区参数，才会导致这两个值丢失。

内容的提问来源于stack exchange，提问作者Charlie Parker