算法面经手撕系列(2)--手撕BatchNormlization

BatchNormlization

  BatchNormlization的编码流程：

1. init阶段初始化$C_{i}n$大小的scale向量和shift向量，同时初始化相同大小的滑动均值向量和滑动标准差向量；
2. forward时沿着非channel维度计算均值、有偏方差
3. 依据得到均值和有偏方差进行归一化
4. 对归一化的结果进行缩放和平移

代码

 代码如下：

class BN(nn.Module):def __init__(self,C_in):super(BN,self).__init__()self.scale=nn.Parameter(torch.ones(C_in).view(1,-1,1,1))self.shift=nn.Parameter(torch.zeros(C_in).view(1,-1,1,1))self.momentum=0.9self.register_buffer('running_mean',torch.zeros(C_in).view(1,-1,1,1))self.register_buffer('running_var',torch.zeros(C_in).view(1,-1,1,1))self.eps=1e-9def forward(self,x):if self.training:N,C,H,W=x.shapemean=x.mean(dim=[0,2,3],keepdim=True)var=x.var(dim=[0,2,3],keepdim=True,unbiased=False)x=(x-mean)/torch.sqrt(var+self.eps)self.running_mean=self.momentum*self.running_mean+(1-self.momentum)*meanself.running_var=self.momentum*self.running_var+(1-self.momentum)*varelse:x=(x-self.running_mean)/torch.sqrt(self.running_var+self.eps)return xif __name__=="__main__":input=torch.rand(10,3,5,5)model=BN(3)res=model(input)print('cool')