论文笔记《BlockDrop: Dynamic Inference Paths in Residual Networks》

论文笔记《BlockDrop: Dynamic Inference Paths in Residual Networks》

paper：https://openaccess.thecvf.com/content_cvpr_2018/html/Wu_BlockDrop_Dynamic_Inference_CVPR_2018_paper.html

code：https://github.com/Tushar-N/blockdrop

Introduction

总体上来说，这篇文章主要目的就是在于希望在推理测试的时候，能够根据输入数据的不同进行不同的路径选择，在ResNet网络基础上可以丢掉一些不需要的block，在保证正确率的水平下减少计算量。

核心就在于BlockDrop，是一种强化学习方法，主要的目的是学习一个policy network，对于一个新图片的输入，通过一个预先定好drop块的resnet网络中输出所有二元决策的后验概率。采用课程学习的方法进行训练，极大程度上提高reward。然后对于训练好的policy network与预训练的resnet进一步进行一个finetune，对于drop功能进行转换，很重要的一点在于drop的决策是一步完成的

BlockDrop将ImageNet上的ResNet-101模型加速了20％，同时保持了相同的76.4％top-1精度，而且可以观察到BlockDrop学习的drop策略与图像中的视觉模式相关。例如，“橙子”类，包含一堆橘子的图像所采用的推理路径与橘子特写图像所采用的推理路径不同。 此外，与包含其他遮挡或背景对象的困难图像相比，具有清晰可见对象的容易图像的BlockDrop策略使用的残留块更少

Approach

对于给定的一张测试图像，目的就是在预训练好的resnet模型中找到block的最佳配置，要求使用的block数量最少，并且不降低精度。采用强化学习中的策略搜索方法来得出最佳block丢弃方案，研究过程中发现resnet非常适合删除，resnet有一个很重要的特点在于在任意两个残差块之间都有直接路径，可以看作是多路径模型。

Policy Network

与标准强化学习不同，我们的策略是一次预测所有的决策，将这个丢弃的策略定义成一个K维的伯努利分布：

[pi_mathbf{W}(mathbf{u|x})=prod_{k=1}^{K}mathbf{s_k^{u_k}}(1-mathbf{s_k})^{1-mathbf{u_k}}\ mathbf{s}=f_{pn}(mathbf{x:W}) ]

(f_{pn})为在权重参数(mathbf{W})下的policy network，(mathbf{s})为sigmoid函数后的输出，向量第k个项(mathbf{s_k})表示其在原始resnet中所对应的残差块被丢弃的概率，基于结果做出一个决策(mathbf{u} in left{0,1 ight}^K)。这里(mathbf{u_k}=0)表示被丢弃，(mathbf{u_k}=1)表示保留，这个决策产生的reward表示为

[R(mathbf{u})= egin{cases} 1-(frac{|mathbf{u}|_0}{K})^2 & mathbb{if correct}\ -gamma & mathbb{otherwise} end{cases} ]

那个分式表示block的利用程度，如果利用率越低则奖励越大，用(gamma)表示惩罚控制准确性与效率，来最大化下面的预期奖励：

[J=mathbb{E}_{mathbf{u}simpi_mathbf{W}}[R(mathbf{u})] ]

总结一下，首先通过模型决定号对于输入的图片用到哪些block，通过这些块向前传播产生一个预测结果，根据正确率和效率生成一个reward

Training the BlockDrop Policy

• 梯度

  首先是梯度，也就是要定义一个loss，为了减少方差，利用了一个self-critical baseline (R(widetilde{u}))

[ abla_mathbf{W}J=mathbb{E}[A abla_mathbf{W}sum_{k=1}^{K}mathbb{log}[mathbf{s_ku_k+(1-s_k)(1-u_k)}]]\ A=R(mathbf{u})-R(mathbf{widetilde{u}}) ]

​ 但是在代码里一直没看懂这一步是怎么搞出来的

​ 对于分布(mathbf{s})，为了防止饱和，用一个参数进行约束(mathbf{s'=alphacdotmathbf{s}+(1-alpha)cdot(1-mathbf{s})})

• 课程学习

  策略梯度方法对于初始化方法非常敏感，由于解空间太大，随机初始化基本都是无效的，我们没有一个所谓的专家示例，为了得到有效的动作序列，我们从课程学习中得到启发，在epoch (t)内，对于(1le t<K)，只学习最后(t)个块的policy，随着(t)增长，更多靠后的块被优化，直到包括所有的块。根据这种方法，首先根据未修改的输入特征优化每个块的激活评估使用程度，随着(t)增加暴露越来越多的不同特征输入，最后(t)个块的policy训练的也越来越好

• 联合finetune

  在课程学习后，policy已经能够决定根据给定的图片原始resnet中哪一个block可以丢弃了。尽管在训练的过程中保证了其正确率，但是移走了一些block还是会不可逆转的导致训练和测试过程中的结果不一致，所以需要联合微调resnet和policy
  

来看一下实际的课程学习相关代码，训练与reward部分，以batchsize 128的cifar数据集resnet110 54个block测试为例，finetune阶段的训练过程其实代码大致相同但是loss还是有点不同

def get_reward(preds, targets, policy):
    # 按batch维度每个值相加 然后除以batch的大小计算平均值
    block_use = policy.sum(1).float()/policy.size(1)
    # 对应了公式里的reward，shape为[128,54]
    sparse_reward = 1.0-block_use**2
    # shape[128,54]
    # 预测结果以及所对应的id,在dim=1的维度上
    _, pred_idx = preds.max(1)
    # 判断是否相等，也就是预测出来结果对不对，data相等则是1 不相等就是0
    match = (pred_idx==targets).data
    # match是一个1 0的[128]
    # 根据match的结果来有不同的reward
    reward = sparse_reward
    #如果不匹配就是惩罚 对应公式
    reward[1-match] = args.penalty
    # 增加一个维度 变成[128,1] 为了对应后面的计算
    reward = reward.unsqueeze(1)
    return reward, match.float()

def train(epoch):
    # Policy网络叫agent train相当于一个开关 启用BN Dropout 与eval在测试时相对应
    agent.train()
    matches, rewards, policies = [], [], []
    for batch_idx, (inputs, targets) in tqdm.tqdm(enumerate(trainloader), total=len(trainloader)):
        inputs, targets = Variable(inputs), Variable(targets).cuda(async=True)
        if not args.parallel:
            inputs = inputs.cuda()
        # cifar: inputs.shape [128, 3, 32, 32]
        probs, value = agent(inputs)
        # print(probs.shape,value.shape) ResNet110 54个block [128, 54] [128, 1]
        #---------------------------------------------------------------------#
        policy_map = probs.data.clone()
        # 在tensor里这么用 list里不行 根据0.5直接判断赋值
        policy_map[policy_map<0.5] = 0.0
        policy_map[policy_map>=0.5] = 1.0
        # tensor不能反向传播，variable可以反向传播
        policy_map = Variable(policy_map)
        # 文中的公式 防止饱和的新公式
        probs = probs*args.alpha + (1-probs)*(1-args.alpha)
        distr = Bernoulli(probs)
        # 按伯努利分布随机取的 根据每个点的概率 取一次0或1
        policy = distr.sample()
        # 处于cl学习状态 只给最后几个学 前面的都不管设置为全1
        if args.cl_step < num_blocks:
            # policymap是模型按照输出来的结果 policy是根据输出的概率值按照伯努利分布生成的结果
            # 除了最后几列还保持原来的 之前的全部为1 只有要训练的那几层还保留
            policy[:, :-args.cl_step] = 1
            policy_map[:, :-args.cl_step] = 1
            # 生成一个全1 mask的作用在于课程学习 保证只学后面那几个layer
            policy_mask = Variable(torch.ones(inputs.size(0), policy.size(1))).cuda()
            # 除了最后那几列其他全部变成0 只学 为1的最后几列
            policy_mask[:, :-args.cl_step] = 0
        else:
            policy_mask = None
        with torch.no_grad():
            v_inputs = Variable(inputs.data)
        # policy_map是输出的预测结果1 0 policy是分布随机采样生成的 带着policy forward
        preds_map = rnet.forward(v_inputs, policy_map)
        preds_sample = rnet.forward(v_inputs, policy)
        reward_map, _ = get_reward(preds_map, targets, policy_map.data)
        reward_sample, match = get_reward(preds_sample, targets, policy.data)
        # 增加一个自监督的过程 来减小方差 参照公式里的A 所以要取两次不同的policy
        advantage = reward_sample - reward_map
        loss = -distr.log_prob(policy)
        # expand_as是变成policy的尺寸 这边loss的计算我就看不懂了 莫非policy就是log后面那一大串 有点像
        loss = loss * Variable(advantage).expand_as(policy)
        if policy_mask is not None:
            loss = policy_mask * loss # mask for curriculum learning
        loss = loss.sum()
        # 限制到下面这个范围内 如果高于上界或低于下界 直接用界来代替
        probs = probs.clamp(1e-15, 1-1e-15)
        # 这里又来了个交叉熵loss beta指的是一个参数entropy multiplier 这里才有点像类似于相对于ground truth增加的loss结果
        entropy_loss = -probs*torch.log(probs)
        entropy_loss = args.beta*entropy_loss.sum()
        loss = (loss - entropy_loss)/inputs.size(0)
        #---------------------------------------------------------------------#
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        matches.append(match.cpu())
        rewards.append(reward_sample.cpu())
        policies.append(policy.data.cpu())
    accuracy, reward, sparsity, variance, policy_set = utils.performance_stats(policies, rewards, matches)
    log_str = 'E: %d | A: %.3f | R: %.2E | S: %.3f | V: %.3f | #: %d'%(epoch, accuracy, reward, sparsity, variance, len(policy_set))
    print (log_str)

Experiment

对于CIFAR数据集，采用了ResNet-32和ResNet-110，分别是15和54个block，每个block有两个卷积层，均匀分成三个stage

对于ImageNet，用ResNet-110，包括33个block，分成4个stage，每个block包括一个三层的bottleneck结构，这些都是pretrained好的SOTA模型

对于Policy网络，将ResNet作为基础模型深度的一小部分，对于CIFAR是一个3个block的ResNet-8，ImageNet是4个block的ResNet-10，此外输入到policy时图片下采样到112x112，所以相对来说这个计算量是可以不用考虑的

随后进行了一系列实验证明了课程学习策略的重要性，联合微调的重要性，与SOTA模型在性能上的对比，只要一步策略的优势，还有语义信息和block丢弃之间的关联。

欢迎批评指正与讨论！关于梯度方面的实现细节，仍有待指点。