剪枝与重参第五课：前置知识

爱听歌的周童鞋

已于 2023-04-13 14:46:19 修改

阅读量1k

点赞数

分类专栏：剪枝与重参文章标签：模型剪枝模型重参数深度学习

于 2023-04-11 23:19:04 首次发布

本博客上原创文章未经本人许可，不得用于商业用途。转载请注明出处，否则保留追究法律责任的权利

本文链接：https://blog.csdn.net/qq_40672115/article/details/130095972

版权

剪枝与重参专栏收录该内容

14 篇文章 85 订阅

订阅专栏

前置知识

前言

手写AI推出的全新模型剪枝与重参课程。记录下个人学习笔记，仅供自己参考。

本次课程主要讲解实战的前置知识。

课程大纲可看下面的思维导图

在这里插入图片描述

1.CIFAR10数据集

1.1 简介

CIFAR10是一个广泛使用的图像分类数据集，由10个类别中的60000张32x32彩色图像组成，每个类别有6000张图像。其中50000张用于训练集，10000张用于测试集。该数据集中的图像均为低分辨率(32x32像素)，使其称为一个小型但具有挑战性的数据集。

CIFAR10数据集的类别包括：airplane(飞机)、automobile(汽车)、bird(鸟)、cat(猫)、deer(鹿)、dog(狗)、frog(青蛙)、horse(马)、ship(船)和truck(卡车)。

下面是CIFAR10数据集中的一些简单示例图：

在这里插入图片描述

1.2 数据集的获取

CIFAR10官网提供了数据集的下载方式，在pytorch中也提供了关于数据集的下载：

import torchvision.datasets as datasets
train_set = datasets.CIFAR10(root='./data.cifar10', train=True, download=True)

运行上述代码后，在data.cifar10文件夹下有下载好的cifar-10-python.tar.gz压缩文件以及解压后的cifar-10-batchs-py文件，该文件夹下的内容有：

batches.meta：包含一个Python字典，包含标签10个类别信息，每个batch的图片数量(10000)，每张图片的像素个数(3072=32x32x3)
data_batch_1 - data_batch_5：训练数据集，每个文件包含10000张图像和对应的标签
readme.html：说明文档
test.batch：测试数据集，包含10000张图像和对应的标签

每个数据文件都是一个二进制文件，可以使用Python的pickle模块进行读取和反序列化。每个数据文件包含一个Python字典，其中data键包含一个10000x3072的NumPy数组，其中每行表示一张图像的颜色通道值(32x32x3=1024个红色像素值+1024个绿色像素值+1024个蓝色像素值)，以及一个label键包含与每个图像相关联的类别标签(0到9之间的整数)。

在读取数据集时，可以使用NumPy数组操作对图像数据进行处理和转换，并使用pickle模块对文件进行反序列化。也可以使用pytorch中的DataLoader对数据集进行加载。

1.3 数据集的加载

CIFAR10数据集加载的示例代码如下：

import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

def create_train_loader(batch_size=16):
    train_set = datasets.CIFAR10(root='./data.cifar10', train=True, download=True)
    mean = train_set.data.mean(axis=(0,1,2)) / 255
    std  = train_set.data.std(axis=(0,1,2))  / 255

    transforms_train = transforms.Compose([
        transforms.Pad(4),
        transforms.RandomCrop(32),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])
    train_set.transform = transforms_train
    train_loader = DataLoader(dataset=train_set, batch_size=batch_size, shuffle=True)
    return train_loader

def create_test_loader(batch_size=16):
    test_set = datasets.CIFAR10(root='./data.cifar10', train=False, download=True)
    mean = test_set.data.mean(axis=(0,1,2)) / 255
    std  = test_set.data.std(axis=(0,1,2))  / 255

    transforms_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])
    test_set.transform = transforms_test
    test_loader = DataLoader(dataset=test_set, batch_size=batch_size, shuffle=True)
    return test_loader

if __name__ == "__main__":
    train_loader = create_train_loader()
    test_loader  = create_test_loader()
    
    import matplotlib.pyplot as plt
    
    # 获取一个batch的数据
    data_iter = iter(test_loader)
    images, labels = data_iter.next()
    
    # 将数据转换为numpy数组
    images = images.numpy()

    # 显示图片
    fig, axes = plt.subplots(nrows=4, ncols=4, figsize=(10,10))

    for i, ax in enumerate(axes.flat):
        # 显示图片
        ax.imshow(images[i].transpose(1,2,0))
        # 设置标签
        ax.set_title(f"Label: {labels[i]}")
        # 隐藏坐标轴
        ax.axis('off')

    plt.show()

上述代码通过PyTorch库的torchvision.datasets来加载CIFAR10数据集，并使用transforms对图像进行预处理，然后使用torch.utils.data中的DataLoader创建数据加载器，最后显示了一个batch的图像数据。

2.VGG网络搭建

2.1 VGGNet

官方网站：Visual Geometry Group Home Page
相关论文：Very Deep Convolutional Networks For Large-scale Image Recognition (2015)

VGGNet是由牛津大学计算机视觉组于2014年提出的一个深度卷积神经网络，它获得了2014年ImageNet图像分类比赛的第二名。VGGNet的特点是采用了非常小的3x3卷积核，使用多个小卷积核来替代大的卷积核，增加网络的深度和非线性表达能力。VGGNet具有很好的可扩展性，可以通过添加更多的卷积层和全连接层来进一步提高网络的性能。VGGNet包含了几个不同深度和宽度的网络结构，其中最有名的是VGG16和VGG19，它们分别包含16和19个卷积层和全连接层。

VGGNet的网络结构非常简单，它包含了若干个卷积层和全连接层，其中卷积层包含了多个卷积核，每个卷积核的大小都是3x3。网络的最后一层是全连接层，用于将卷积层的输出映射到类别标签上。VGGNet的网络结构中，每个卷积层都采用了相同的结构，即两个3x3的卷积核，每个卷积核后面都跟了一个ReLU激活函数，最后是一个2x2的最大池化层。这个基本单元被称为VGG块，网络中的所有卷积层都由多个VGG块组成。

下图是VGG-16的网络结构图：

在这里插入图片描述

不同深度和宽度的VGG网络结构的配置信息如下：

defaultcfg = {
    11 : [64,     'M', 128,      'M', 256, 256,           'M', 512, 512,           'M', 512, 512          ],
    13 : [64, 64, 'M', 128, 128, 'M', 256, 256,           'M', 512, 512,           'M', 512, 512          ],
    16 : [64, 64, 'M', 128, 128, 'M', 256, 256, 256,      'M', 512, 512, 512,      'M', 512, 512, 512     ],
    19 : [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512],
}

2.2 VGG网络实现

VGG网络搭建的示例代码如下：

import torch.nn as nn

defaultcfg = {
    11 : [64,     'M', 128,      'M', 256, 256,           'M', 512, 512,           'M', 512, 512          ],
    13 : [64, 64, 'M', 128, 128, 'M', 256, 256,           'M', 512, 512,           'M', 512, 512          ],
    16 : [64, 64, 'M', 128, 128, 'M', 256, 256, 256,      'M', 512, 512, 512,      'M', 512, 512, 512     ],
    19 : [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512],
}

class VGG(nn.Module):
    def __init__(self, num_classes=10, depth=11, cfg=None):
        super().__init__()
        if not cfg:
            cfg = defaultcfg[depth]
        
        self.features = self.make_layers(cfg)
        self.classifier = nn.Linear(cfg[-1], num_classes)

    def make_layers(self, cfg):
        layers = []
        in_channels = 3
        for l in cfg:
            if l == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                conv2d = nn.Conv2d(in_channels, l, kernel_size=3, padding=1, bias=False)
                layers += [conv2d, nn.BatchNorm2d(l), nn.ReLU(inplace=True)]
                in_channels = l
        return nn.Sequential(*layers)
    
    def forward(self, x):
        x = self.features(x)
        x = nn.AvgPool2d(2)(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

if __name__ == "__main__":
    vgg = VGG()
    print(vgg)

我们搭建的VGG-11网络与原始的VGG网络不同，在features后采取的是AvgPool2d进行池化且为了简单化只使用了一个全连接层。

3.Batch Normalize

3.1 简介

相关论文：Batch Normalization- Accelerating Deep Network Training b y Reducing Internal Covariate Shift (2015)
知乎解读：Batch Normalization原理与实战

Batch Normalize(批标准化)是一种深度神经网络中常用的正则化方法，旨在缓解深度神经网络中梯度消失或梯度爆炸的问题，加速训练过程并提高模型的性能。

Batch Normalize在训练过程中，对每个 minibatch 的输出进行标准化，即对每个特征在 batch 维度上进行标准化，使得输出的均值和标准差分别为 0 和 1。这样做的好处在于，使得每个层的输入都是以相同的方式进行标准化，从而加速了训练过程。

具体来说，Batch Normalize 可以分为以下几个步骤：

对于输入特征 $x$ ，计算其均值 $\mu$ 和标准差 $\sigma$ 。
标准化：将特征 $x$ 标准化为 $\hat{x} = \frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}$ ，其中 $\epsilon$ 是一个小的常数，防止除数为 0。
对标准化后的特征进行缩放和平移： $\gamma\hat{x} + \beta$ ，其中 $\gamma$ 和 $\beta$ 是可学习的参数，使得模型可以自适应地选择适当的缩放和平移，从而提高模型的拟合能力。
对于每个 minibatch，通过梯度下降来更新 $\gamma$ 和 $\beta$ 。

Batch Normalize在深度神经网络中广泛使用，可以使模型训练更加稳定和快速，从而提高模型的准确性。

3.2 BN层实现

BN层搭建的示例代码如下：


import torch
import torch.nn as nn

class BatchNorm(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum

        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta  = nn.Parameter(torch.zeros(num_features))

        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))

    def forward(self, x):
        if self.training:
            mean = x.mean(dim=(0,2,3), keepdim=True)
            var  = x.var(dim=(0,2,3), keepdim=True)

            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.squeeze()
            self.running_var  = (1 - self.momentum) * self.running_var  + self.momentum * var.squeeze()
        else:
            mean = self.running_mean.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
            var  = self.running_var.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)

        x = (x - mean) / torch.sqrt(var + self.eps)

        x = x * self.gamma.unsqueeze(-1).unsqueeze(-1) + self.beta.unsqueeze(-1).unsqueeze(-1)

        return x

上面的示例代码用于实现Batch Normalize，主要有以下几点说明：

由于self.gamma和self.beta都是可学习的参数，其值需要在训练过程中被更新，因此它们被定义为nn.Parameter对象。nn.Parameter是Tensor的子类，它的主要作用是为了将一个Tensor封装成一个 Parameter对象。这样做的好处是，将一个Tensor封装成Parameter对象后，该Tensor会被自动注册为模型的参数，可以被自动更新。
register_buffer是nn.Module类中的一个方法，它用于注册一个持久化的buffer，该buffer不需要梯度，且在调用to()方法时会自动将其移动到相应的设备上。在Batch Normalization中，running_mean和running_var是在训练过程中不断更新的均值和方差，它们需要在每次前向传播时被保存下来。因此，将它们注册为buffer可以保证它们被自动保存和移动到正确的设备上，而且不会被当做模型参数进行优化。
对于CNN而言，输入的数据一般是4D的张量即(batch,channels,height,weight)，对于每个channel，需要对batch个样本求均值和方差，所以求取mean和var是(0,2,3)。至于keepdim=True 的含义是指在求取均值和方差时是否保持维度不变。如果keepdim=True，则均值和方差张量的维度与输入张量维度相同，否则在求均值和方差时会进行降维操作。在Batch Normalization中，keepdim=True 是为了保证均值和方差张量的维度与gamma和beta张量的维度相同，从而能够进行后续的运算。
running_mean和running_var的计算方式是对每个Batch的均值和方差进行动量平均。在新的Batch到来时，运用动量平均，将原有running_mean和新的均值进行一定比例的加权平均，以此来逐步调整整个数据集的分布，从而更好地适应新的数据分布。这样做的目的是在训练过程中更好地适应不同的数据分布，从而提高网络的泛化能力。其中动量momentum为0.1是较为常见的选择
squeeze()表示将tensor维度为1的维度去掉。在BatchNorm的实现中，mean和var计算得到的是形状为(1,C,1,1)的tensor，其中C为特征的通道数。使用squeeze()可以将tensor的形状变为(C,)，方便后续计算。unsqueeze()是PyTorch中用于增加维度的方法，它的作用是在指定的维度上增加一个维度，其参数是增加的维度的索引。

4.L1&L2正则

我们所说的正则化，就是在原来的Loss Function的基础上，加上了一些正则化项或者称为模型复杂度惩罚项

Loss Function
$\frac{1}{N} *\sum\limits^{N}_{i=1}(y_i - w^Tx_i)^2$
假设 $L (w)$ 在0处的导数为0，即达到最优解：
$\frac{\partial L(w)}{\partial w}\Bigm|_{w=0} = d = 0$

4.1 L1正则化(Lasso回归)

加上L1正则项(Lasso 回归)： $C||w||_1$
损失函数：

$L_{L1}(w)= L(w) + \lambda|w|$
导数：

$\frac{\partial L(w)}{\partial w} \Bigm|_{w=0^-} = d - \lambda$

$\frac{\partial L(w)}{\partial w} \Bigm|_{w=0^+} = d + \lambda$

在这里插入图片描述

在剪枝中，L1正则化会用在Batch Normalization上面的 $\gamma$ 上面，实现稀疏训练

4.2 L2正则化(岭回归)

加上L2正则项(岭回归)： $C||w||^2_2$
损失函数：

$L_{L2}(w) = L(w)+ \lambda w^2$
导数：

$\frac{\partial L(w)}{\partial w} \Bigm|_{w=0} = d + 2\lambda w =0$

在这里插入图片描述

4.3 思考

问题1：为什么使用L1正则化可以实现稀疏训练？

L1正则化是一种对模型权重进行惩罚的方法，它将权重中的小值变为0，从而实现模型的稀疏化。在Batch Normalization中， $\gamma$ 是缩放因子，它用于缩放归一化的输出，而且 $\gamma$ 的初始值通常被设置为1，如果对其进行L1正则化惩罚，会使得模型更倾向于将一些通道的权重设置为0，从而实现通道的剪枝，减少模型参数量和计算量。因此，在剪枝中，L1正则化被广泛应用于Batch Normalization的缩放因子 $\gamma$ 上。(from chatGPT)

问题2：为什么使用L1正则化不使用L2正则化？

L1正则化在想要寻找一个能够大幅减少权值数量的最优解时很有用。L1正则化对权值施加的惩罚不像L2正则化那样平滑，它倾向于让一些权值变为0。对于 $\gamma$ 系数，因为它们的值用于调节每个通道的缩放因子，使其接近于1，而一些通道的重要性可能不如其他通道。因此，使用L1正则化有助于找到仅仅使用少量通道可以获得相同性能的 $\gamma$ 系数。而L2正则化倾向于使得所有 $\gamma$ 系数都很小但非零。因此，在使用L1正则化时，可以通过稀疏化权重获得一些模型压缩和加速的好处。(from chatGPT)

5.train

模型稀疏训练具体实现流程

5.1 parse_opt

利用argparse命令行参数解析模块传入模型训练时的参数，示例代码如下：

import argparse

def parse_opt():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch Slimming CIFAR training')
    parser.add_argument('--dataset', type=str, default='cifar100', help='training dataset (default: cifar100)')
    parser.add_argument('--sparsity-regularization', '-sr', dest='sr', action='store_true', help='train with channel sparsity regularization')
    parser.add_argument('--s', type=float, default=0.0001, help='scale sparse rate (default: 0.0001)')
    parser.add_argument('--refine', default='', type=str, metavar='PATH', help='path to the pruned model to be fine tuned')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=256, metavar='N', help='input batch size for testing (default: 256)')
    parser.add_argument('--epochs', type=int, default=160, metavar='N', help='number of epochs to train (default: 160)')
    parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)')
    parser.add_argument('--lr', type=float, default=0.1, metavar='LR', help='learning rate (default: 0.1)')
    parser.add_argument('--momentum', type=float, default=0.9, metavar='M', help='SGD momentum (default: 0.9)')
    parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)')
    parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')
    parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=100, metavar='N', help='how many batches to wait before logging training status')
    parser.add_argument('--save', default='./logs', type=str, metavar='PATH', help='path to save prune model (default: current directory)')
    parser.add_argument('--arch', default='vgg', type=str,  help='architecture to use')
    parser.add_argument('--depth', default=19, type=int, help='depth of the neural network')

    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_opt()
    print(args)

5.2 train

训练函数实现的示例代码如下：

import argparse
from models import VGG
from utils import create_train_loader

import torch
import torch.optim as optim
import torch.nn.functional as F
import torch.nn as nn

def parse_opt():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch Slimming CIFAR training')
    parser.add_argument('--dataset', type=str, default='cifar100', help='training dataset (default: cifar100)')
    parser.add_argument('--sparsity-regularization', '-sr', dest='sr', action='store_true', help='train with channel sparsity regularization')
    parser.add_argument('--s', type=float, default=0.0001, help='scale sparse rate (default: 0.0001)')
    parser.add_argument('--refine', default='', type=str, metavar='PATH', help='path to the pruned model to be fine tuned')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=256, metavar='N', help='input batch size for testing (default: 256)')
    parser.add_argument('--epochs', type=int, default=160, metavar='N', help='number of epochs to train (default: 160)')
    parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)')
    parser.add_argument('--lr', type=float, default=0.1, metavar='LR', help='learning rate (default: 0.1)')
    parser.add_argument('--momentum', type=float, default=0.9, metavar='M', help='SGD momentum (default: 0.9)')
    parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)')
    parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')
    parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=100, metavar='N', help='how many batches to wait before logging training status')
    parser.add_argument('--save', default='./logs', type=str, metavar='PATH', help='path to save prune model (default: current directory)')
    parser.add_argument('--arch', default='vgg', type=str,  help='architecture to use')
    parser.add_argument('--depth', default=19, type=int, help='depth of the neural network')

    args = parser.parse_args()
    return args


def updateBN():
    # 更新Batch Normalization中的gamma参数，使用L1正则化来实现稀疏训练
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.weight.grad.data.add_(args.s * torch.sign(m.weight.data))

def train(epoch):
    model.train()   # 将模型设置训练模式
    for batch_idx, (data, target) in enumerate(train_loader):
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()   # 梯度清零
        output = model(data)    # 前向传播
        loss = F.cross_entropy(output, target)
        loss.backward() #反向传播
        if args.sr:
            updateBN()  # 稀疏训练，更新BN层的gamma参数
        optimizer.step()    # 参数更新
        if batch_idx % args.log_interval == 0:  # 消息打印
            print('Train Epoch: {} [{}/{} ({:.1f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), 
                len(train_loader.dataset),
                100. * batch_idx / len(train_loader), 
                loss.item()))   


if __name__ == '__main__':
    args = parse_opt()
   
    args.cuda = not args.no_cuda and torch.cuda.is_available()

    model = VGG()
    if args.cuda:
        model.cuda()
    train_loader = create_train_loader()

    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    train(epoch=1)

5.3 test

测试函数实现的示例代码如下：

def test():
    model.eval()    # 将模型设置为测试模式
    test_loss = 0
    correct = 0
    with torch.no_grad():   # 不进行梯度计算
        for data, target in test_loader:
            if args.cuda:
                data, target = data.cuda(), target.cuda()
            output = model(data)    # 前向传播得到预测结果
            test_loss = F.cross_entropy(output, target, reduction='sum').item() # 计算损失
            pred = output.data.max(1, keepdim=True)[1]  # 获取预测类别
            correct += pred.eq(target.data.view_as(pred)).cpu().sum()   # 统计正确的数量
    
    test_loss /= len(test_loader.dataset)   # 计算平均损失
    # 打印测试结果
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.1f}%)\n'.format(
        test_loss, correct, 
        len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
    return correct / float(len(test_loader.dataset))    # 返回测试准确率

5.4 save_checkpoint

保存模型的示例代码如下：

import os
import shutil

def save_checkpoint(state, is_best, filepath):
    torch.save(state, os.path.join(filepath, 'checkpoint.pth'))
    if is_best:
        shutil.copyfile(os.path.join(filepath, 'checkpoint.pth'), os.path.join(filepath, 'model_best.pth'))

5.5 完善示例代码

完整的train.py示例代码如下：

import os
import torch
import argparse

import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F


from models.vgg import VGG
from utils import get_training_dataloader, get_test_dataloader, save_checkpoint


def parse_opt():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch Slimming CIFAR training')
    parser.add_argument('--dataset', type=str, default='cifar10', help='training dataset (default: cifar100)')
    parser.add_argument('--sparsity-regularization', '-sr', dest='sr', action='store_true', help='train with channel sparsity regularization')
    parser.add_argument('--s', type=float, default=0.0001, help='scale sparse rate (default: 0.0001)')
    parser.add_argument('--refine', default='', type=str, metavar='PATH', help='path to the pruned model to be fine tuned')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=256, metavar='N', help='input batch size for testing (default: 256)')
    parser.add_argument('--epochs', type=int, default=160, metavar='N', help='number of epochs to train (default: 160)')
    parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)')
    parser.add_argument('--lr', type=float, default=0.1, metavar='LR', help='learning rate (default: 0.1)')
    parser.add_argument('--momentum', type=float, default=0.9, metavar='M', help='SGD momentum (default: 0.9)')
    parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)')
    parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')
    parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=100, metavar='N', help='how many batches to wait before logging training status')
    parser.add_argument('--save', default='./logs', type=str, metavar='PATH', help='path to save prune model (default: current directory)')
    parser.add_argument('--arch', default='vgg', type=str,  help='architecture to use')
    parser.add_argument('--depth', default=19, type=int, help='depth of the neural network')

    args = parser.parse_args()
    return args

# additional subgradient descent on the sparsity-induced penalty term
def updateBN():
    for m in model.modules():
        # Check if the module is a BatchNorm2d layer
        if isinstance(m, nn.BatchNorm2d):
            # Calculate the L1 regularization term and add it to the weight gradients
            # args.s is a scalar value that determines the strength of the regularization
            # torch.sign(m.weight.data) returns the sign of the weight parameters
            m.weight.grad.data.add_(args.s*torch.sign(m.weight.data))  # L1


def train(epoch):
    # Set the model to training mode
    model.train()
    # Loop through the batches in the training data
    for batch_idx, (data, target) in enumerate(train_loader):
        # Move the data and target tensors to the GPU if args.cuda is True
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        # Zero out the gradients in the optimizer
        optimizer.zero_grad()
        # Forward pass: compute the output of the model on the input data
        output = model(data)
        # Compute the loss between the output and target labels
        loss = F.cross_entropy(output, target)
        # Backward pass: compute the gradients of the loss w.r.t. the model parameters
        loss.backward()
        # If args.sr is True, apply L1 regularization to the Batch Normalization layers
        if args.sr:
            updateBN()
        # Update the model parameters using the optimizer
        optimizer.step()
        # Print the training loss and progress at regular intervals
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.1f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), 
                len(train_loader.dataset),
                100. * batch_idx / len(train_loader), 
                loss.item()))


def test():
    # Set the model to evaluation mode
    model.eval()
    # Initialize test loss and correct predictions
    test_loss = 0
    correct = 0
    # Turn off gradient calculation during inference
    with torch.no_grad():    
        # Loop through the test data
        for data, target in test_loader:
            # Move the data and target tensors to the GPU if args.cuda is True
            if args.cuda:
                data, target = data.cuda(), target.cuda()
            # Compute the output of the model on the input data
            output = model(data)
            # Compute the test loss and add it to the running total
            test_loss += F.cross_entropy(output, target, reduction='sum').item()
            # Compute the predictions from the output using the argmax operation
            pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
            # Compute the number of correct predictions and add it to the running total
            correct += pred.eq(target.data.view_as(pred)).cpu().sum()
    # Compute the average test loss and accuracy
    test_loss /= len(test_loader.dataset)
    # Print the test results
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.1f}%)\n'.format(
        test_loss, correct, 
        len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
    return correct / float(len(test_loader.dataset))


if __name__ == '__main__':
    # Parse command line arguments
    args = parse_opt()
    # Check if CUDA is available and set args.cuda flag accordingly
    args.cuda = not args.no_cuda and torch.cuda.is_available()
    # Set the random seed for PyTorch and CUDA if args.cuda is True
    torch.manual_seed(args.seed)
    if args.cuda:
        torch.cuda.manual_seed(args.seed)
    # Create the save directory if it does not exist
    if not os.path.exists(args.save):
        os.makedirs(args.save)
        
    # Set kwargs to num_workers=1 and pin_memory=True if args.cuda is True, 
    # otherwise kwargs is an empty dictionary
    kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
    
    # Create data loaders for the CIFAR10 dataset 
    # using the get_training_dataloader() and get_test_dataloader() functions
    if args.dataset == 'cifar10':
        train_loader = get_training_dataloader(batch_size=args.batch_size, **kwargs)
        test_loader  = get_test_dataloader(batch_size=args.test_batch_size, **kwargs)
        
    # Load a pre-trained VGG model if args.refine is not None, 
    # otherwise create a new VGG model
    if args.refine:
        checkpoint = torch.load(args.refine)
        model = VGG(depth=args.depth, cfg=checkpoint['cfg'])
        model.load_state_dict(checkpoint['state_dict'])
    else:
        model = VGG(depth=args.depth)
    # Move the model to the GPU if args.cuda is True
    if args.cuda:
        model.cuda()
    # Set up the optimizer with Stochastic Gradient Descent (SGD) 
    # and the specified learning rate, momentum, and weight decay
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)

    if args.resume:
        # Check if the checkpoint file exists
        if os.path.isfile(args.resume):
            # If the checkpoint file exists, print a message indicating that it's being loaded
            print("=> loading checkpoint '{}'".format(args.resume))
            # Load the checkpoint file
            checkpoint = torch.load(args.resume)
            # Update the start epoch and best precision variables from the checkpoint
            args.start_epoch = checkpoint['epoch']
            best_prec1 = checkpoint['best_prec1']
            
            # Load the model state dictionary and optimizer state dictionary from the checkpoint
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            # Print a message indicating that the checkpoint has been loaded
            print("=> loaded checkpoint '{}' (epoch {}) Prec1: {:f}"
                .format(args.resume, checkpoint['epoch'], best_prec1))
        else:
            # If the checkpoint file does not exist, print an error message
            print("=> no checkpoint found at '{}'".format(args.resume))
            
    # Initialize the best test accuracy to 0
    best_prec1 = 0.
    # Loop through the epochs, starting from args.start_epoch and continuing until args.epochs
    for epoch in range(args.start_epoch, args.epochs):
        # If the current epoch is at 50% or 75% of the total epochs, 
        # reduce the learning rate by a factor of 10
        if epoch in [args.epochs*0.5, args.epochs*0.75]:
            for param_group in optimizer.param_groups:
                param_group['lr'] *= 0.1
        # Train the model on the training data for the current epoch
        train(epoch)
        # Evaluate the model on the test data and compute the top-1 test accuracy
        prec1 = test()
        # Check if the current test accuracy is better than the previous best test accuracy
        is_best = prec1 > best_prec1
        # Update the best test accuracy and save a checkpoint of the model and optimizer state
        best_prec1 = max(prec1, best_prec1)
        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'best_prec1': best_prec1,
            'optimizer': optimizer.state_dict(),
        }, is_best, filepath=args.save)
    # Print the best test accuracy achieved during training
    print("Best accuracy: "+str(best_prec1))