【pytorch】Batch-norm layer loss计算

Batchnor-layer loss calculation

在实现data-free KD 时用到的loss函数的计算。这个函数放在batch norm 1d层之后计算不同层输出之间的差异。

思路：通过pytorch设置hook钩子函数，获取某一个特定层的输出；

可以计算指定的层的输出之间的差异，即输出的统计量的差异，一般以均值和方差来衡量这个统计量；

重点：输出是什么，shape或者是size是什么样子的？在内存中是怎么存储的？

下面以一个例子来看一下具体的维度的变化，given a three-dimensional tensor, [batch,text_length,embedding feature num]

import torch

def Gaussian_kd(mean, var, T_mean, T_var):
    num = (mean-T_mean)**2 + var
    denom = 2*T_var
    std = torch.sqrt(var)
    T_std = torch.sqrt(T_var)

    return num/denom - torch.log(std/T_std) - 0.5  # 对应着论文中的公式(5),计算的是两个高斯分布之间额KL divergence

x = torch.randn(2, 1, 5) # 2个batch，单词长度为1，embedding之后的feature num是1500
nch = x.shape[1] # 1,获得句子的长度
mean = x.mean([0,2]) # 除了句子长度之外，获得batch以及embeding feature维度上的均值以及方差

var = x.permute(1,0,2).contiguous().view([nch,-1]).var(1,unbiased=False) # reshape，只获得通道维度上的统计量的值，为了内存连续读取，使用contiguous以及view函数

y = torch.randn(2, 1, 5) 
nch_y = y.shape[1]
mean_y = y.mean([0,2])
var_y = y.permute(1,0,2).contiguous().view([nch_y,-1]).var(1,unbiased=False)

g_loss = Gaussian_kd(mean,var,mean_y,var_y)