Dynamic Slimmable Network-(CVPR21-ORAL)综合分析

文章提出了一种动态剪枝策略，通过动态宽度可变超网络（Dynamic Slimmable Supernet），解决了传统剪枝方法硬件实际加速效果很低的问题。并提出了动态宽度门控（Dynamic Slimming Gate）对网络进行瘦身。
下面就文章原理和代码来综合分析：

动态宽度可变超网络

动态宽度可变网络（DS-Net）通过学习一个宽度可变超网络和一个动态门控机制来实现不同样本的动态路由。如上图所示，DS-Net中的超网络（上图黄色框）是指承担主要任务的整个模块。相比之下，动态门控（上图蓝色框）是一系列预测模块，它们将输入样本路由到超网络的不同宽度的子网络。

1.动态超网络（supernet）和动态可切分（slice-able）卷积
为避免产生稀疏channel，作者提出动态可切分（slice-able）卷积，通过预测出的剪枝率 ，动态的切分使用卷积的前 × n个滤波器（n为总滤波器数）。通过堆叠动态可切分（slice-able）卷积并禁用动态门控，就形成了类似slimmablenetwork的动态超网络。
如下为conv2d的slice方法。

    def forward(self, x):
        if self.prev_channel_choice is None:
            self.prev_channel_choice = self.channel_choice
        if self.mode == 'dynamic' and isinstance(self.channel_choice, tuple):
            weight = self.weight
            if not self.in_chn_static:
                if isinstance(self.prev_channel_choice, int):
                    self.running_inc = self.in_channels_list[self.prev_channel_choice]
                    weight = self.weight[:, :self.running_inc]
                else:
                    self.running_inc = torch.matmul(self.prev_channel_choice[0], self.in_channels_list_tensor)
            if not self.out_chn_static:
                self.running_outc = torch.matmul(self.channel_choice[0], self.out_channels_list_tensor)

            output = F.conv2d(x,
                              weight,
                              self.bias,
                              self.stride,
                              self.padding,
                              self.dilation,
                              self.groups)
            if not self.out_chn_static:
                output = apply_differentiable_gate_channel(output,
                                                           self.channel_choice[0],
                                                           self.out_channels_list)
            self.prev_channel_choice = None
            self.channel_choice = -1
            return output
        else:
            if not self.in_chn_static:
                self.running_inc = x.size(1)
            if not self.out_chn_static:
                self.running_outc = self.out_channels_list[self.channel_choice]
            weight = self.weight[:self.running_outc, :self.running_inc]
            bias = self.bias[:self.running_outc] if self.bias is not None else None
            self.running_groups = 1 if self.groups == 1 else self.running_outc
            self.prev_channel_choice = None
            self.channel_choice = -1
            return F.conv2d(x,
                            weight,
                            bias,
                            self.stride,
                            self.padding,
                            self.dilation,
                            self.running_groups)

2.In-place Ensemble Bootstrapping（IEB）
之前的slimmable network训练使用in-place distillation方法：最宽的子网络学习预测真实标签，同时生成软标签，并通过知识蒸馏的方式来训练其他较窄的子网络。但in-place distillation训练很不稳定，权重在训练早期会大幅突变，并可能导致模型最终训练失败或性能损失。

为此，作者提出In-place Ensemble Bootstrapping（IEB）策略来稳定动态超网络的训练并最终提高模型性能。首先，使用超网络的滑动平均（EMA）网络来生成训练子网络的软标签，因为EMA网络提供的目标更加稳定和精准。其次，使用包括最宽子网络和随机宽度子网络的多个模型的概率集成（probability ensemble）作为训练最窄网络的目标，因为多模型集成可以提供更多样、更泛化、更精准的软标签。（见上图）
一下代码为IEB策略代码。

    for batch_idx, (input, target) in enumerate(loader):
        last_batch = batch_idx == last_idx
        data_time_m.update(time.time() - end)
        if not args.prefetcher:
            input, target = input.cuda(), target.cuda()

        sample_list = ['largest', 'uniform', 'uniform', 'smallest']
        guide_list = []

        for sample_idx, model_mode in enumerate(sample_list):
            seed = seed + 1
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            if hasattr(model, 'module'):
                model.module.set_mode(model_mode)
            else:
                model.set_mode(model_mode)
            #model.blocks._modules['3'].first_block.conv_dw.running_inc running_outc
            output = model(input)
 
            if model_mode == 'largest':
                loss = loss_fn(output, target)#原始标签计算loss
                if args.ieb:
                    with torch.no_grad():
                        if hasattr(model_ema.ema, 'module'):
                            model_ema.ema.module.set_mode(model_mode)
                        else:
                            model_ema.ema.set_mode(model_mode)
                        output_largest = model_ema.ema(input)
                    guide_list.append(output_largest)
                loss_largest = loss
            elif model_mode != 'smallest':
                if args.ieb:
                    loss = distill_loss_fn(output, F.softmax(output_largest, dim=1))
                    # with torch.no_grad():
                    #     guide_output = model_ema.ema(input)
                    with torch.no_grad():
                        if hasattr(model_ema.ema, 'module'):
                            model_ema.ema.module.set_mode(model_mode)
                        else:
                            model_ema.ema.set_mode(model_mode)
                        guide_output = model_ema.ema(input)
                    guide_list.append((guide_output))
                    # guide_list.append((output.detach()))
                else:
                    loss = loss_fn(output, target)
            else:  # 'smallest'
                soft_labels_ = [torch.unsqueeze(guide_list[idx], dim=2) for
                                idx in range(len(guide_list))]
                soft_labels_softmax = [F.softmax(i, dim=1) for i in soft_labels_]
                soft_labels_softmax = torch.cat(soft_labels_softmax, dim=2).mean(dim=2)
                if args.ieb:
                    loss = distill_loss_fn(output, soft_labels_softmax)
                else:
                    loss = loss_fn(output, target)
                loss_smallest = loss

            loss = loss / optimizer_step

            if use_amp:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

动态宽度门控

作者设计了双头（double headed）动态宽度门控的结构，并提出gate的训练策略sandwich gate sparsification（SGS）。

首先，作者提出通过当前层输入特征来预测出一个one-hot向量，对应选择剪枝率列表中的一个值。输入特征首先通过全局池化（global pooling）来消除空间维度，再通过两个以ReLU相隔的全连接层，并求argmax，得到one-hot向量。由于这一结构与频道注意力相似，作者将两者除最后一层全连接外的层进行共享，形成包含动态宽度头和频道注意力头的双头动态宽度门控。
以下为gate的code

class MultiHeadGate(nn.Module):
    def __init__(self, in_chs, se_ratio=0.25, reduced_base_chs=None, act_layer=nn.ReLU,
                 attn_act_fn=sigmoid, divisor=1, channel_gate_num=None, gate_num_features=1024):
        super(MultiHeadGate, self).__init__()
        self.attn_act_fn = attn_act_fn
        self.channel_gate_num = channel_gate_num
        reduced_chs = make_divisible((reduced_base_chs or in_chs[-1]) * se_ratio, divisor)
        self.avg_pool = DSAdaptiveAvgPool2d(1, channel_list=in_chs)
        self.conv_reduce = DSpwConv2d(in_chs, [reduced_chs], bias=True)
        self.act1 = act_layer(inplace=True)
        self.conv_expand = DSpwConv2d([reduced_chs], in_chs, bias=True)

        self.has_gate = False
        if channel_gate_num > 1:
            self.has_gate = True
            # self.gate = nn.Sequential(DSpwConv2d([reduced_chs], [gate_num_features], bias=True),
            #                           act_layer(inplace=True),
            #                           nn.Dropout2d(p=0.2),
            #                           DSpwConv2d([gate_num_features], [channel_gate_num], bias=True))
            self.gate = nn.Sequential(DSpwConv2d([reduced_chs], [channel_gate_num], bias=False))

        self.mode = 'largest'
        self.keep_gate, self.print_gate, self.print_idx = None, None, None
        self.channel_choice = None
        self.initialized = False
        if self.attn_act_fn == 'tanh':
            nn.init.zeros_(self.conv_expand.weight)
            nn.init.zeros_(self.conv_expand.bias)

    def forward(self, x):
        x_pool = self.avg_pool(x)
        x_reduced = self.conv_reduce(x_pool)
        x_reduced = self.act1(x_reduced)
        attn = self.conv_expand(x_reduced)
        if self.attn_act_fn == 'tanh':
            attn = (1 + attn.tanh())
        else:
            attn = self.attn_act_fn(attn)
        x = x * attn

        if self.mode == 'dynamic' and self.has_gate:
            channel_choice = self.gate(x_reduced).squeeze(-1).squeeze(-1)
            self.keep_gate, self.print_gate, self.print_idx = gumbel_softmax(channel_choice, dim=1, training=self.training)
            self.channel_choice = self.print_gate, self.print_idx
        else:
            self.channel_choice = None

        return x

    def get_gate(self):
        return self.channel_choice

2. Sandwich Gate Sparsification（SGS）

由于argmax不可导，之前使用argmax作为网络中间层的工作一般使用gumbel-softmax作为替代，来近似求导，以便梯度回传。但是，本文作者发现采用这个方法进行gate训练时，很容易使其gate坍塌成静态。

为此，作者提出Sandwich Gate Sparsification训练策略。首先，每个输入样本都按，“是否能被最窄网络正确预测”，分为容易和困难两类。然后，将两类难易样本打上one-hot标签，使用交叉熵直接优化gate。这种训练方式避免了间接和近似的梯度回传，克服了gate收敛困难的问题，并提高了gate的动态多样性。

    for batch_idx, (input, target) in enumerate(loader):
        last_batch = batch_idx == last_idx
        data_time_m.update(time.time() - end)
        if not args.prefetcher:
            input, target = input.cuda(), target.cuda()
            #torch.jit.trace(model, input).save('dsnet1.pt')

        if last_batch or (batch_idx + 1) % optimizer_step == 0:
            optimizer.zero_grad()
        # generate online labels
        with torch.no_grad():
            set_model_mode(model, 'smallest')
            output = model(input)
            conf_s, correct_s = accuracy(output, target, no_reduce=True)
            gate_target = [torch.LongTensor([0]) if correct_s[0][idx] else torch.LongTensor([3])
                           for idx in range(correct_s[0].size(0))]
            gate_target = torch.stack(gate_target).squeeze(-1).cuda()
        # =============

        set_model_mode(model, 'dynamic')
      
        output = model(input)
      )

        if hasattr(model, 'module'):
            model_ = model.module
        else:
            model_ = model

        #  SGS Loss
        gate_loss = 0
        gate_num = 0
        gate_loss_l = []
        gate_acc_l = []
        for n, m in model_.named_modules():
            if isinstance(m, MultiHeadGate):
                if getattr(m, 'keep_gate', None) is not None:
                    gate_num += 1
                    g_loss = loss_fn(m.keep_gate, gate_target)
                    gate_loss += g_loss
                    gate_loss_l.append(g_loss)
                    gate_acc_l.append(accuracy(m.keep_gate, gate_target, topk=(1,))[0])

        gate_loss /= gate_num

        #  MAdds Loss
        running_flops = add_flops(model)
        if isinstance(running_flops, torch.Tensor):
            running_flops = running_flops.float().mean().cuda()
        else:
            running_flops = torch.FloatTensor([running_flops]).cuda()


        flops_loss = (running_flops / 1e9) ** 2

        #  Target Loss, back-propagate through gumbel-softmax
        ce_loss = loss_fn(output, target)

        loss = gate_loss + ce_loss + 0.5 * flops_loss
        # loss = ce_loss
        acc1 = accuracy(output, target, topk=(1,))[0]

        if use_amp:
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            loss.backward()

评估gate

    for batch_idx, (input, target) in enumerate(loader):
        last_batch = batch_idx == last_idx
        data_time_m.update(time.time() - end)
        if not args.prefetcher:
            input, target = input.cuda(), target.cuda()
        # generate online labels
        with torch.no_grad():
            set_model_mode(model, 'smallest')
            output = model(input)
            conf_s, correct_s = accuracy(output, target, no_reduce=True)
            gate_target = [torch.LongTensor([0]) if correct_s[0][idx] else torch.LongTensor([3])
                           for idx in range(correct_s[0].size(0))]
            gate_target = torch.stack(gate_target).squeeze(-1).cuda()
        # =============
        set_model_mode(model, 'dynamic')
        output = model(input)

        if hasattr(model, 'module'):
            model_ = model.module
        else:
            model_ = model

        gate_acc_l = []
        for n, m in model_.named_modules():
            if isinstance(m, MultiHeadGate):
                if getattr(m, 'keep_gate', None) is not None:
                    gate_acc_l.append(accuracy(m.keep_gate, gate_target, topk=(1,))[0])

        running_flops = add_flops(model)
        if isinstance(running_flops, torch.Tensor):
            running_flops = running_flops.float().mean().cuda()
        else:
            running_flops = torch.FloatTensor([running_flops]).cuda()

        loss = loss_fn(output, target)
        prec1, prec5 = accuracy(output, target, topk=(1, 5))

评估validate_slim

    for choice in range(args.num_choice):
        eval_metrics.append(validate_slim(model,
                                          loader_eval,
                                          validate_loss_fn,
                                          args,
                                          model_mode=choice))

validate_slim

    with torch.no_grad():
        for batch_idx, (input, target) in enumerate(loader):
            if not isinstance(model_mode, str):
                if hasattr(model, 'module'):
                    model.module.set_mode('uniform', choice=model_mode)
                else:
                    model.set_mode('uniform', choice=model_mode)
            else:
                if hasattr(model, 'module'):
                    model.module.set_mode(model_mode)
                else:
                    model.set_mode(model_mode)
            last_batch = batch_idx == last_idx
            if not args.prefetcher:
                input = input.cuda()
                target = target.cuda()

            output = model(input)
            running_flops = add_flops(model)
            if isinstance(running_flops, torch.Tensor):
                running_flops = running_flops.float().mean().cuda()
            else:
                running_flops = torch.FloatTensor([running_flops]).cuda()

            # augmentation reduction
            reduce_factor = args.tta
            if reduce_factor > 1:
                output = output.unfold(0, reduce_factor, reduce_factor).mean(dim=2)
                target = target[0:target.size(0):reduce_factor]

            loss = loss_fn(output, target)
            prec1, prec5 = accuracy(output, target, topk=(1, 5))

可以看出，评估gate网络用的是dynamic模式；对网络进行测试时，四个模型一起预测评估。
需要注意的是：该代码用来保留4个channels选项。

    for n, m in model.named_modules():
        if len(getattr(m, 'in_channels_list', [])) > 4:
            m.in_channels_list = m.in_channels_list[start_chn_idx:4]
            m.in_channels_list_tensor = torch.from_numpy(
                np.array(m.in_channels_list)).float().cuda()
        if len(getattr(m, 'out_channels_list', [])) > 4:
            m.out_channels_list = m.out_channels_list[start_chn_idx:4]
            m.out_channels_list_tensor = torch.from_numpy(
                np.array(m.out_channels_list)).float().cuda()

总结

从代码中可以看出，（Stage I: Supernet Training）通过对 [‘largest’, ‘uniform’, ‘uniform’, ‘smallest’]四个submodel进行软标签式训练，该阶段不训练gate，只训练supernet。 （Stage II: Gate Training）该阶段在训练出原始模型后，通过SGS来训练gate（gate在stage的3block.first_block.gate.gate，输出为[66,4]）。最终在前传阶段，通过将mode置为dynamic，使得模型进行weight*gate形式前传来预测最终label。对于检测网络，感觉不易实现该压缩算法。

参考文章

• https://zhuanlan.zhihu.com/p/370208935
• https://www.163.com/dy/article/G6QJKLHQ0511CQLG.html
• https://zhuanlan.zhihu.com/p/354043252
• CVPR 2021 ORAL