【论文阅读】Long-Tailed Recognition via Weight Balancing（CVPR2022）附MaxNorm的代码

论文

问题：真实世界中普遍存在长尾识别问题，朴素训练产生的模型在更高准确率方面偏向于普通类，导致稀有的类别准确率偏低。
key:解决LTR的关键是平衡各方面，包括数据分布、训练损失和学习中的梯度。
文章主要讨论了三种方法： L2normalization, weight decay, and MaxNorm
本文提出了一个两阶段训练的范式：
a. 利用调节权重衰减的交叉熵损失学习特征。
b. 通过调节权重衰减和Max Norm使用类平衡损失学习分类器。
一些有用的看法：

研究表明，与联合训练特征学习和分类器学习的模型相比，解耦特征学习和分类器学习导致了显著的改进。
根据基准测试结果，通过集成专家模型或采用主动数据增强技术的自监督预训练来实现最好精度。
研究发现，SGD动量导致LTR出现问题，阻碍了进一步改善。
最近，Kang等人令人信服地证明了阶段性训练对LTR很重要。
权重衰减有助于学习隐藏层的平衡权重。
重要的是，我们的探索发现，虽然在分类器上使用L2规范化约束进行训练比简单训练有所改进，但它的表现不如下面描述的其他两个正则化。
与严格将所有滤波器权重的范数值设置为1的L2归一化不同，MaxNorm放松了这一约束，允许权重在训练期间在范数球内移动。
权重衰减中，不同数据集的最优λ各不相同——较大的数据集需要较小的权重衰减，直观地说，因为在更多数据上学习有助于泛化，因此需要较少的正则化。
单阶段使用不平衡损失训练效果不好的原因：虽然他们没有解释为什么具有类平衡损失的单阶段训练表现不佳，但直观地说，这是因为类平衡损失人为地放大了从罕见的类训练数据计算的梯度，这损害了特征表示学习，从而损害了最终的LTR性能。
本文作者使用了weight decay和max norm两种方法结合，因为发现两个结合效果更好。让模型不同类之间权重相差不会很大的同时，还能让这些权重缓慢增加。
下面这幅图就是解释了这些方法的特点。

第一个就是普通方法训练的，它常见的类别权重增长快。
第二个是L2 normalization，它把所有类别的权重都限定在一个常数。
第三个是权重衰减，它的所有类的权重小，而且权重在增长。
第四个是MaxNorm，它限制最大的权重。
第五个是权重衰减和MaxNorm，会导致范数中的权重较小且平衡。

使用方法

weight decay

先定义好权重衰减的值。

weight_decay = 0.1 #weight decay value

然后在优化器中调用。Adam还有其他的都有weight_decay。

optimizer = optim.SGD([{'params': active_layers, 'lr': base_lr}], lr=base_lr, momentum=0.9, weight_decay=weight_decay)

MaxNorm

就是这个论文中的regularizers.py中的代码。只要会使用就好。就是要是不是作者代码中的模型的话，model.encoder.fc还需要根据自己的代码修改。

#使用前先定义好初始化好
pgdFunc = MaxNorm_via_PGD(thresh=thresh)
pgdFunc.setPerLayerThresh(model) # set per-layer thresholds这个是计算模型每一层的权重的阈值，这篇论文中只计算最后线性层的权重，并对最后线性层的权重进行限制

当模型训练一个epoch结束后，对已经更新完毕的模型权重进行限制，如果超过阈值就进行更新，让权重在最大范数的约束下。

 if pgdFunc:# Projected Gradient Descent
     pgdFunc.PGD(model)#对权重进行限制

import torch
import torch.nn as nn
import math
# The classes below wrap core functions to impose weight regurlarization constraints in training or finetuning a network.

class MaxNorm_via_PGD():
    def __init__(self, thresh=1.0, LpNorm=1, tau=1):
        self.thresh = thresh
        self.LpNorm = LpNorm
        self.tau = tau
        self.perLayerThresh = []

    def setPerLayerThresh(self, model):#根据指定的模型设置每层的阈值
        #set pre-layer thresholds
        self.perLayerThresh = []

        for curLayer in [model.encoder.fc.weight, model.encoder.fc.bias]:#遍历模型的最后两层
            curparam = curLayer.data#获取当前层的数据
            if len(curparam.shape) <= 1:#如果层只有一个维度，是一个偏置或者是一个1D的向量，则设置这一层的阈值为无穷大，继续下一层
                self.perLayerThresh.append(float('inf'))
                continue
            curparam_vec = curparam.reshape((curparam.shape[0], -1))#如果不是，把权重张量展开
            neuronNorm_curparam = torch.linalg.norm(curparam_vec, ord=self.LpNorm, dim=1).detach().unsqueeze(-1)#沿着第一维计算P番薯，结果存储
            curLayerThresh = neuronNorm_curparam.min() + self.thresh*(neuronNorm_curparam.max() - neuronNorm_curparam.min())#计算每一层的阈值及神经元范数的最小值加上最大值和最小值之间的缩放差
            self.perLayerThresh.append(curLayerThresh)#每层阈值存储

    def PGD(self, model):#定义PGD函数，用于在模型的参数上执行投影梯度下降，试试最大范数约束
        if len(self.perLayerThresh) == 0:#如果每层的阈值是空，用setPerLayerThresh方法初始化
            self.setPerLayerThresh(model)
        for i, curLayer in enumerate([model.encoder.fc.weight, model.encoder.fc.bias]):#遍历模型的最后两层
            curparam = curLayer.data#获取当前层的数据张量值
            curparam_vec = curparam.reshape((curparam.shape[0], -1))#变成一维
            neuronNorm_curparam = (torch.linalg.norm(curparam_vec, ord=self.LpNorm, dim=1)**self.tau).detach().unsqueeze(-1)#在最后加一维
            #计算权重张量中每行神经元番薯的tau次方
            scalingVect = torch.ones_like(curparam)#创建一个形状与当前层数据相同的张量，用1初始化
            curLayerThresh = self.perLayerThresh[i]#获取阈值

            idx = neuronNorm_curparam > curLayerThresh#创建bool保存超过阈值的神经元
            idx = idx.squeeze()#
            tmp = curLayerThresh / (neuronNorm_curparam[idx].squeeze())**(self.tau)#根据每层的阈值和超过阈值的神经元番薯计算缩放因子
            for _ in range(len(scalingVect.shape)-1):#扩展缩放因子以匹配当前层数据的维度
                tmp = tmp.unsqueeze(-1)

            scalingVect[idx] = torch.mul(scalingVect[idx],tmp)
            curparam[idx] = scalingVect[idx] * curparam[idx]
            curparam[idx] = scalingVect[idx] * curparam[idx]#通过缩放值更新当前层的数据，以便对超过阈值的神经元进行缩放。完成权重更新

如果使用原来的代码报错的可以看下面这个

我的网络只有一层是线性层idx = idx.squeeze()，idx是(1,1)形状的，squeeze就没了，所以报错，如果有这个原因的可以改成idx = idx.squeeze(1)。maxnorm只改最后两层/一层权重所以，定义了一个列表存储线性层只取最后两层或者一层。

class MaxNorm_via_PGD():
    # learning a max-norm constrainted network via projected gradient descent (PGD)
    def __init__(self, thresh=1.0, LpNorm=2, tau=1):
        self.thresh = thresh
        self.LpNorm = LpNorm
        self.tau = tau
        self.perLayerThresh = []

    def setPerLayerThresh(self, model):
        # set per-layer thresholds
        self.perLayerThresh = []#存储每一层的阈值
        self.last_two_linear_layers = []#提取线性层
        for name, module in model.named_children():
            if isinstance(module, nn.Linear):
                self.last_two_linear_layers.append(module)

        for linear_layer in self.last_two_linear_layers[-min(2, len(self.last_two_linear_layers)):]:  # here we only apply MaxNorm over the last two layers
            curparam = linear_layer.weight.data
            if len(curparam.shape) <= 1:
                self.perLayerThresh.append(float('inf'))
                continue
            curparam_vec = curparam.reshape((curparam.shape[0], -1))
            neuronNorm_curparam = torch.linalg.norm(curparam_vec, ord=self.LpNorm, dim=1).detach().unsqueeze(-1)
            curLayerThresh = neuronNorm_curparam.min() + self.thresh * (
                        neuronNorm_curparam.max() - neuronNorm_curparam.min())
            self.perLayerThresh.append(curLayerThresh)

    def PGD(self, model):
        if len(self.perLayerThresh) == 0:
            self.setPerLayerThresh(model)
        for i, curLayer in enumerate([self.last_two_linear_layers[-min(2,
                                                             len(self.last_two_linear_layers))]]):  # here we only apply MaxNorm over the last two layers

            curparam = curLayer.weight.data

            curparam_vec = curparam.reshape((curparam.shape[0], -1))
            neuronNorm_curparam = (
                        torch.linalg.norm(curparam_vec, ord=self.LpNorm, dim=1) ** self.tau).detach().unsqueeze(-1)
            scalingVect = torch.ones_like(curparam)
            curLayerThresh = self.perLayerThresh[i]

            idx = neuronNorm_curparam > curLayerThresh
            idx = idx.squeeze(1)
            tmp = curLayerThresh / (neuronNorm_curparam[idx].squeeze()) ** (self.tau)
            for _ in range(len(scalingVect.shape) - 1):
                tmp = tmp.unsqueeze(-1)

            scalingVect[idx] = torch.mul(scalingVect[idx], tmp)
            curparam[idx] = scalingVect[idx] * curparam[idx]