論文筆記《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\right\}^K\)。這裡\(\mathbf{u_k}=0\)表示被丟棄，\(\mathbf{u_k}=1\)表示保留，這個決策產生的reward表示為

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

那個分式表示block的利用程度，如果利用率越低則獎勵越大，用\(\gamma\)表示懲罰控制準確性與效率，來最大化下面的預期獎勵：

\[J=\mathbb{E}_{\mathbf{u}\sim\pi_\mathbf{W}}[R(\mathbf{u})] \]

總結一下，首先通過模型決定號對於輸入的圖片用到哪些block，通過這些塊向前傳播產生一個預測結果，根據正確率和效率生成一個reward

Training the BlockDrop Policy

• 梯度

  首先是梯度，也就是要定義一個loss，為了減少方差，利用了一個self-critical baseline \(R(\widetilde{u})\)

\[\nabla_\mathbf{W}J=\mathbb{E}[A\nabla_\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'=\alpha\cdot\mathbf{s}+(1-\alpha)\cdot(1-\mathbf{s})}\)

• 課程學習

  策略梯度方法對於初始化方法非常敏感，由於解空間太大，隨機初始化基本都是無效的，我們沒有一個所謂的專家示例，為了得到有效的動作序列，我們從課程學習中得到啟發，在epoch \(t\)內，對於\(1\le 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丟棄之間的關聯。

歡迎批評指正與討論！關於梯度方面的實現細節，仍有待指點。