Caffe的BN参数转移到PyTorch上

Caffe的模型参数转移到PyTorch上，最需要注意的是Caffe中的BN层+Scale层与PyTorch中BN层区别，理解之后BN层参数转移很简单。

也可以参考GitHub代码：pytorch-caffe

Batch Normalization

 

Caffe的BN+Scale等于PyTorch的BN，二者都可以用上述公式表达。

1. 在Caffe中，BN层负责计算均值E和方差Var；Scale层进行归一化，包含两个参数λ和β。
   BN层参数:

 

message BatchNormParameter {
  // 当为真，使用保存的均值和方差，否则使用滑动平均计算新的方差和均值
  optional bool use_global_stats = 1;
  //滑动平均的系数
  optional float moving_average_fraction = 2 [default = .999];
  // 平滑，防止除以0
  optional float eps = 3 [default = 1e-5];
}

BN层有三个参数保存：均值E、方差Var、滑动系数；Scale层保存两个参数：λ和β。

1. 在PyTorch中，BN层定义：

 

class _BatchNorm(Module):
# eps防止除以0
# momentum 滑动平均的系数(实际代码中等价于caffe中的 0.999)
# affine是否归一化
# track_running_stats等价于Caffe中的use_global_stats
    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
                 track_running_stats=True):

BN层中保存四个参数：均值running_mean、方差running_var、权重weight、偏置bias。

BN层参数从Caffe转换到PyTorch

发现Caffe中的BN层参数有5个，而PyTorch中只有四个，多出一个滑动系数。参考pytorch-caffe的做法，是把Caffe中的均值和方差先除以滑动系数，再转到PyTorch中。例如：

 

注意点

• 在Caffe中训练时设置use_global_stats=false，测试时需要改为use_global_stats=true。
• 在PyTorch中model.train()默认track_running_stats=true,测试时model. eval()默认track_running_stats=false。