剪枝论文三（ Sparse Structure Selection）

本文介绍一种剪枝方法（ Sparse Structure Selection）。

1. 核心思想

如图为论文中提出的网络框架。F表示残差函数。灰色的block、group和neuron表示它们是不活跃的，对应的比例因子为0，可以被修剪。

根据作者在论文中所描述的，改论文的贡献体现在以下2个方面:

1. 作者提出了一个统一的cnn模型训练和剪枝框架。特别地，通过在cnn的某些结构上引入缩放因子和相应的稀疏正则化，将其转化为一个联合稀疏正则化优化问题。
2. 作者利用改进的随机加速近端梯度(APG)方法，利用稀疏正则化联合优化cnn和标度因子的权重。与以往采用启发式方法来强制稀疏性的方法相比，该方法无需进行微调和多阶段优化，具有更稳定的收敛性和更好的结果。

2. 损失函数

作者引入了一种新的参数——缩放因子$\lambda$来缩放某些特定结构(神经元、组或块)的输出，并在训练过程中增加了$\lambda$的稀疏性约束。目标是得到一个稀疏的$\lambda$。即，如果${\lambda }_i=0$，则可以安全地删除相应的结构，因为它的输出对后续的计算没有贡献。

作者给出的目标函数为：

式中，$L(y_i,C(x_i,W,\lambda ))$是样本$x_i$上的损失
$R(\cdot )$是应用于每个权重的非结构化正则化，例如$l_2$范数作为权重衰减
$R_s(\cdot )$为$\lambda$带权$\gamma$的稀疏正则化。在论文中使用了最常用的凸松弛$l_1$范数，定义为 $\gamma ||\lambda |{|}_1$。

3. 代码实现

3.1 优化器

1. 使用SGD优化权重weight_decay
2. 使用APGNAG优化gamma

# 使用SGD优化权重
optimizer1 = SGD([{'params': model.parameters()}], lr=args.lr,
                 momentum=args.momentum, weight_decay=args.weight_decay)

# 如果剪枝，使用APGNAG优化gamma
if args.sss:
    optimizer2 = APGNAG([{'params': model.lambda_block}],
                        lr=args.lr, momentum=args.momentum, gamma=args.gamma)

3.2 训练

根据以下损失函数学习：

式中，

1. $L(y_i,C(x_i,W,\lambda ))$是样本$x_i$上的损失，使用交叉熵损失。
2. 代码中未使用
3. $R_s(\cdot )$为$\gamma ||\lambda |{|}_1$，使用l1_loss损失。

def train(epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data), Variable(target)

        # 如果剪枝
        if args.sss:
            optimizer1.zero_grad()
            optimizer2.zero_grad()

            # 将model中的前向传递函数foward函数替换为foward_sss
            output = model.forward_sss(data)
            loss = F.cross_entropy(output, target)
            pred = output.data.max(1, keepdim=True)[1]
            loss.backward()

            # 额外增加了r1_loss，用来优化lambda_block
            r1_loss = args.gamma * F.l1_loss(model.lambda_block, torch.zeros(model.lambda_block.size()).cuda(), reduction='sum')
            r1_loss.backward()
            optimizer1.step()
            optimizer2.step()

        # 不剪枝，只使用SGD正常优化权重
        else:
            optimizer1.zero_grad()
            output = model(data)
            loss = F.cross_entropy(output, target)
            pred = output.data.max(1, keepdim=True)[1]
            loss.backward()
            optimizer1.step()

3.4 前向传递函数

训练中采用的前向传递函数由原有的foward替换为foward_sss函数

class ResNet_cifar(nn.Module):
    def forward_sss(self, x):
        x = self.conv1(x)
        num_block = 0
        for block in self.layer1:
            x = block.forward_sss(x, self.lambda_block[num_block])
            num_block += 1
        for block in self.layer2:
            x = block.forward_sss(x, self.lambda_block[num_block])
            num_block += 1
        for block in self.layer3:
            x = block.forward_sss(x, self.lambda_block[num_block])
            num_block += 1
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

其中block.forward_sss相比于block.forward的变化只有一个。就是增加：

output *= lambda

block.forward_sss如下

    def forward_sss(self, x, drop):
        '''
        剪枝：对比forward，增加了out *= drop
        argument:
            drop: lambda_block[num_block]，第num_block个模块中使用的lambda
        '''
        residual = x
        if self.downsample is not None:
            residual = self.downsample(x)

        out = self.bn1(x)
        out = self.relu(out)
        out = self.conv1(out)

        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv2(out)

        out = self.bn3(out)
        out = self.relu(out)
        out = self.conv3(out)

        out *= drop # 增加的代码
        out += residual

        return out

3.5 学习率更新

两个优化器的学习率都随着训练的增加而降低：

for epoch in range(args.start_epoch, args.epochs):
    if epoch in [args.epochs*0.5, args.epochs*0.75]:
        # 学习率随着训练的增加而降低
        for param_group in optimizer1.param_groups:
            param_group['lr'] *= 0.1
        if args.sss:
            for param_group in optimizer2.param_groups:
                param_group['lr'] *= 0.1

源代码