模型压缩——如何进行神经网络架构搜索？

1. 引言

前面剪枝的思路是对一个已经存在的模型做参数量的压缩，那我们不禁会想一个问题：既然模型存在可压缩的空间，那为什么不从一开始就找到这个冗余少的模型，并在这个小参数模型上进行训练呢？

这里的困难点在于，要想构建一个有效模型，需要满足延时低、存储少、消耗少，同时还要保持模型精度等多个目标。但是手工设计模型结构的方式，并没有一个有效的手段能够同时在多个目标间达到平衡。所以我们就会想，能不能通过自动化的方式来尝试多种组合，自动选择出最优目标？

神经网络架构搜索(Neural Architecture Search, NAS)，可以用来解决这个问题，它把神经网络的结构、参数量、组合方式都当做搜索空间，利用某种搜索算法对所有可能的参数和结构组合进行性能评估，以此来快速的搜索到合适的神经网络架构。

它的基本思路是：先规定一个搜索空间，通过给定的搜索策略，得到与策略对应的神经网络架构模型，通过数据集评估该模型的性能，再将该结果反馈给搜索策略，以决定是否得到了更优的神经网络架构。不断迭代尝试新的策略，最终得到所有策略中最优的神经网络架构。

接下来我们就通过实际操作来演示一下如何进行神经网络架构搜索。

2. 模型和数据

首先，我们引进模型剪枝中已经编写过的数据、训练、评估相关的代码。

%run lenet.py

os.environ['CUDA_VISIBLE_DEVICES'] = "2" 

2.1 模型结构改造

在模型剪枝一文中介绍了经典的卷积神经网络模型LeNet, 但它的模型结构是内部固定的，我们需要一个能通过外部策略动态配置内部结构的模型类。

因此，我们需要对LeNet进行改造，将卷积层的关键参数（如卷积核大小、通道数），以及全连接层的关键参数（如out_features)以可配置的方式定义。

class LeNetPro(nn.Module):
    def __init__(self, conv1_channel=6, conv1_kernel=5, conv2_channel=16, conv2_kernel=5, fc1_size=128, fc2_size=84):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=conv1_channel, kernel_size=conv1_kernel)
        self.conv2 = nn.Conv2d(in_channels=conv1_channel, out_channels=conv2_channel, kernel_size=conv2_kernel)
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        linear_input_dim = conv2_channel * (((28 - conv1_kernel + 1)//2 - conv2_kernel + 1)//2) **2
        self.fc1 = nn.Linear(in_features=linear_input_dim, out_features=fc1_size)
        self.fc2 = nn.Linear(in_features=fc1_size, out_features=fc2_size)
        self.fc3 = nn.Linear(in_features=fc2_size, out_features=10)
        
    def forward(self, x):
        x = self.maxpool(F.relu(self.conv1(x)))
        x = self.maxpool(F.relu(self.conv2(x)))
        
        x = x.view(x.size()[0], -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

经过这个调整后，我们就可以在创建模型实例时通过参数指定模型内部结构。

model = LeNetPro(conv1_channel=6, conv2_channel=16, fc1_size=128, fc2_size=84)
model

    LeNetPro(
      (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
      (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
      (maxpool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (fc1): Linear(in_features=256, out_features=128, bias=True)
      (fc2): Linear(in_features=128, out_features=84, bias=True)
      (fc3): Linear(in_features=84, out_features=10, bias=True)
    )

2.2 加载数据

数据部分，我们将模型剪枝一节加载数据的过程封装为一个load_data方法，直接加载得到训练集和测试集。

train_loader, test_loader = load_data("./data", batch_size=64)

取一些样例可视化内容显示：

inputs, targets = next(iter(train_loader))
samples = [(inputs[i], targets[i]) for i in range(len(inputs))]
show_image(samples)

3. 搜索空间

进行网络架构搜索之前需要先定义一个搜索空间，告知搜索算法可以尝试哪些网络结构的配置，架构搜索算法将会根据这些配置来生成不同的模型架构。一般有以下原则：

• 多性样：搜索空间应反应网络架构的各种可能性，包含不同类型的层、通道数、卷积核尺寸，使得算法有可探索的空间。
• 有效性：每种类型的参数应在合理的范围内递增或递减，例如：卷积核尺寸3、5、7。
• 实用性：搜索空间不应过于狭窄，以确保算法能探索到潜在的最佳架构，同时也不应过于庞大，以免影响搜索效率。
• 兼容性：不同层的设置应考虑层间的兼容性，例如下一层的输入通道数应等于前一层的输出通道数。

我们先定义以下搜索空间。

conv1_channel_list = [3, 6, 9]  # 卷积层一的通道数选项
conv1_kernel_list = [3, 5, 7]   # 卷积层一的卷积核尺寸选项
conv2_channel_list = [12, 16, 20]  # 卷积层二的通道数选项
conv2_kernel_list = [3, 5, 7]   # 卷积层二的卷积核尺寸选项
fc1_size_list = [64, 128, 256]  # 全连接层一的节点数选项
fc2_size_list = [32, 84, 120]   # 全连接层二的节点数选项

上面每个超参数都定义了三个选项，这6个超参数的每种选项组合都代表一种神经网络结构，可以通过itertools.product方法得到所有可能的超参数组合。

itertools.product(*iterables, repeat=1)：函数用于生成可迭代对象的笛卡尔积，它允许输入多个集合，它会对这些集合中的每一个元素进行组合，输出即所有可能的组合结果。

import itertools
import random

combinations = list(itertools.product(conv1_channel_list, conv1_kernel_list, conv2_channel_list, conv2_kernel_list, fc1_size_list, fc2_size_list))
len(combinations)

    (729,

共有729组合，这些组合就是所有我们能找到的神经网络结构。

4. 性能评估

定义出搜索空间后，我们需要对每种参数组合对应的模型进行训练和性能评估，目的是找到准确率最高的模型和参数组合。

模型训练超参数配置：

lr = 0.01        # 学习率
momentum = 0.5   # 冲量，用于加速收敛
num_epochs = 5    # 训练的轮数
loss_fn = nn.CrossEntropyLoss()  # 交叉熵损失
device = 'cuda' if torch.cuda.is_available() else 'cpu'

定义模型的训练评测方法，目的是找到指定模型在训练过程中表现出的最优性能数据。

注：每个模型在train_loader指定的数据集上训练num_epochs轮。每轮训练完后进行性能评测，所有轮训练完后，选择性能最好的checkpoint作为最终模型状态。

def train(num_epochs, model, train_loader, test_loader, loss_fn, optimizer, device):
    best_accuracy = 0 
    best_checkpoint_state = None
    for i in range(num_epochs):
        train_epoch(model, train_loader, loss_fn, optimizer, device)
        accuracy = evaluate(model, test_loader, device)
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            best_checkpoint_state = copy.deepcopy(model.state_dict())
        print(f'Epoch {i + 1:>2d} Accuracy {accuracy:.2f}% / Best Accuracy: {best_accuracy:.2f}%')
    model.load_state_dict(best_checkpoint_state)
    model_accuracy = evaluate(model, test_loader, device)
    print(f"model accuracy: {model_accuracy:.2f}%")
    return model_accuracy

5. 搜索策略

搜索策略相当于学习方法，它研究的是通过什么学习方法能够快速、准确的找到最优解。

这里采用最简单的方法：随机搜索法，先将搜索空间中的所有组合随机打乱顺序。

注：随机搜索法的特点是简单，但它需要对每一种模型结构进行训练，所以组合数量多时很耗时。业界有一种one-shot搜索算法，先设计一种超网络（包含所有子结构），只训练一次，所有子结构便可以直接从超网络获得其权重，无需从头训练。但本文未作研究，有兴趣可以参考文章最后链接。

random.shuffle(combinations)
combinations[:10]

 [(9, 7, 20, 7, 256, 84),
  (6, 7, 16, 7, 256, 84),
  (6, 5, 16, 3, 128, 84),
  (3, 7, 20, 7, 128, 32),
  (3, 7, 20, 7, 64, 84),
  (3, 7, 16, 7, 128, 84),
  (3, 7, 20, 3, 128, 120),
  (3, 5, 16, 5, 128, 84),
  (3, 5, 12, 3, 64, 32),
  (3, 7, 16, 7, 128, 32)])

对所有组合依次遍历，创建对应的模型实例并评测其性能表现，并不断与已经评测过的模型进行比较，记录性能表现最好的模型状态和超参组合。

%%time
overall_best_model = None
overall_best_accuracy = 0
overall_best_model_info = ""
overall_records = []

print("Starting search...")

for item in combinations:
    model = LeNetPro(*item).to(device)
    num_params = count_parameters(model)
    param_info = f"conv1_channel:{item[0]}, conv1_kernel:{item[1]}, "\
        f"conv2_channel:{item[2]}, conv2_kernel:{item[3]}, "\
        f"fc1_size:{item[4]}, fc2_size:{item[5]}"
    print("Testing configuration: ", param_info)
    print(f"model parameters: {num_params}")
        
    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)
    accuracy = train(num_epochs, model, train_loader, test_loader, loss_fn, optimizer, device)
    if accuracy > overall_best_accuracy:
        overall_best_model = model
        overall_best_accuracy = accuracy
        overall_best_model_info = param_info
    overall_records.append((accuracy, num_params, item))

print("Completed search.")
print(f"Best model configuration: {overall_best_model_info}")
print(f"Best model accuracy: {overall_best_accuracy:.2f}%")

执行上述代码后开始搜索，搜索过程截图示意如下：

整个搜索的过程耗时很长，搜索完成后输出如下信息：

Completed search.
Best model configuration: conv1_channel:9, conv1_kernel:5, conv2_channel:16, conv2_kernel:7, fc1_size:256, fc2_size:120
Best model accuracy: 98.56%

6.搜索结果分析

上面在搜索过程中，记录了所有模型结构的评测记录，这里将所有模型的记录按照准确率排序。

overall_records.sort(key=lambda x: x[0], reverse=True)  

以生成Markdown表格的形式输出性能排名前10的参数组合。

from IPython.display import Markdown

top_n = 10
markdown_datas = []  

# 添加表头  
markdown_datas.append("| Rank | Parameters | Accuracy (%) | conv1_channel | conv1_kernel | conv2_channel | conv2_kernel | fc1 | fc2 |")  
markdown_datas.append("|------|------------|--------------|---------------|--------------|---------------|--------------|-----|-----|")  

# 添加数据行  
for rank, (accuracy, num_params, info) in enumerate(overall_records[:top_n], start=1):  
    markdown_datas.append(f"| {rank} | {num_params} | {accuracy:.2f} |{info[0]}|{info[1]}|{info[2]}|{info[3]}|{info[4]}|{info[5]}|")  

# 输出整个Markdown表格  
final_markdown_table = "\n".join(markdown_datas)  
Markdown(final_markdown_table) 

Rank	Parameters	Accuracy (%)	conv1_channel	conv1_kernel	conv2_channel	conv2_kernel	fc1	fc2
1	76476	98.56	9	5	16	7	256	120
2	135184	98.51	3	3	16	3	256	120
3	29492	98.51	9	5	12	3	64	120
4	82400	98.48	3	3	12	5	256	120
5	64050	98.43	6	5	16	3	128	84
6	60908	98.42	9	5	12	5	256	32
7	151854	98.37	6	3	20	3	256	84
8	33520	98.32	9	5	16	3	64	84
9	74046	98.32	6	5	16	7	256	120
10	77070	98.31	6	5	20	3	128	84

从上面的结果来看，大多数排名靠前的模型具有如下特征：

1. 在参数数量和准确率之间表现出一定的平衡，参数规模在6万-8万之间时容易达到较好的性能。
2. 较小的卷积核出现的频率更高（通常3x3和5x5），小卷积核可能更有助于捕捉局部特征。
3. 大多数高准确率的模型在全连接层输出的维度较高，全连接层增大神经元数量可能有助于模型能力的增强。

小结：本节介绍了神经网络架构搜索的基本操作步骤和相关准备工作，先对模型结构进行了可配置化改造以适应架构搜索的需要，再定义可搜索的组合空间，设计性能评估方法，最后执行搜索策略。同时，在搜索过程中收集每个模型结构的性能数据和参数组合，以便最后选择在性能和参数规模方面表现都比较好的模型。

深度学习终结了手工设计特征的时代，而NAS——模型搜索（architecture search）的最终目标是：终结人工设计架构。

相关阅读

• 模型剪枝
• 模型压缩概览
• One-shot模型搜索