基于NAS的模型压缩自动化让算法自己搜索最优子网络结构一、手工压缩的不可扩展性传统模型压缩流程参见0717/8.md讨论的移动端压缩方案高度依赖人工经验选择哪些层剪枝、哪些通道保留、各层的量化位宽如何分配——这些决策由工程师基于直觉和试错完成。当模型架构变复杂如包含非标准算子的Transformer变体、部署约束增多同时满足延迟、内存、能耗三项约束时手工搜索的效率急剧下降。神经架构搜索Neural Architecture Search, NAS在模型压缩中的应用——即压缩感知的NASCompression-Aware NAS——将找到满足约束的最优压缩方案建模为一个自动化搜索问题。搜索空间从原始NAS的选择操作类型卷积/池化/跳跃连接转变为选择每个算子的压缩策略剪枝比例/量化位宽/是否蒸馏目标函数从单纯的最大化精度变为在满足延迟/内存约束下最大化精度。flowchart TB A[预训练大模型br/教师] -- B{NAS压缩搜索} B -- C[定义搜索空间] C -- C1[每层的剪枝比例br/{0%, 10%, ..., 90%}] C -- C2[每层的量化位宽br/{2, 4, 8, 32}] C -- C3[每层的kernel大小br/{1, 3, 5}] B -- D[定义目标函数] D -- D1[精度: val_accuracy] D -- D2[约束: latency 10ms] D -- D3[约束: model_size 50MB] B -- E[选择搜索策略] E -- E1[强化学习 (RL)] E -- E2[进化算法 (EA)] E -- E3[可微分搜索 (DARTS式)] C -- F[搜索循环] D -- F E -- F F -- G{满足约束?} G --|否| F G --|是| H[最优压缩方案] H -- I[微调恢复精度]二、压缩NAS的搜索空间设计压缩NAS的搜索空间由每个可压缩算子的压缩动作组成。以卷积网络为例一个典型的搜索空间包括宽度维度每层保留的通道比例如[0.25, 0.5, 0.75, 1.0]深度维度是否跳过特定层0或1精度维度每层的权重和激活位宽如[2, 4, 8]结构维度每层的卷积核大小如[1, 3, 5]搜索空间的规模是所有维度选项数的乘积——以ResNet-50为例约50个可压缩层每层4个通道选项理论搜索空间为4^50 ≈ 10^30直接穷举不可行。import torch import torch.nn as nn from typing import List, Dict, Tuple, Optional import numpy as np import copy class CompressionSearchSpace: 压缩NAS的搜索空间定义。 将模型压缩决策编码为可搜索的离散动作序列。 def __init__(self, model: nn.Module): self.model model self.compressible_layers self._identify_compressible_layers() # 定义每层的候选压缩动作 self.action_space { layer_name: { prune_ratio: [0.0, 0.1, 0.3, 0.5, 0.7], # 剪枝比例 quant_bits: [2, 4, 8, 32], # 量化位宽(32不量化) skip: [0, 1], # 是否跳过该层 } for layer_name in self.compressible_layers } def _identify_compressible_layers(self) - List[str]: 识别模型中可压缩的层。 只考虑Conv2d和Linear层含可训练权重。 BN层和激活函数不直接压缩。 compressible [] for name, module in self.model.named_modules(): if isinstance(module, (nn.Conv2d, nn.Linear)): compressible.append(name) return compressible def sample_random_architecture(self) - Dict: 从搜索空间中随机采样一个压缩方案。 Returns: {layer_name: {prune_ratio: x, quant_bits: y, ...}} architecture {} for layer_name, actions in self.action_space.items(): architecture[layer_name] { action: np.random.choice(options) for action, options in actions.items() } return architecture def get_space_size(self) - int: 计算搜索空间的理论大小组合数。 total 1 for actions in self.action_space.values(): layer_options 1 for options in actions.values(): layer_options * len(options) total * layer_options return total def apply_compression_scheme( model: nn.Module, compression_scheme: Dict ) - Tuple[nn.Module, Dict]: 将NAS搜索到的压缩方案应用到模型上。 根据搜索结果的各层配置 1. 执行结构化剪枝按prune_ratio裁剪通道 2. 执行量化按quant_bits 3. 按skip标志移除层 Args: model: 原始FP32模型 compression_scheme: {layer: {prune_ratio, quant_bits, skip}} Returns: (压缩后的模型, 压缩统计信息) compressed_model copy.deepcopy(model) stats {original_params: sum(p.numel() for p in model.parameters())} for name, module in compressed_model.named_modules(): if name not in compression_scheme: continue config compression_scheme[name] # 应用剪枝 if config.get(prune_ratio, 0) 0 and hasattr(module, weight): prune_ratio config[prune_ratio] weight module.weight.data # 按L1范数剪枝输出通道 if isinstance(module, (nn.Conv2d, nn.Linear)): importance weight.abs().sum(dim(1, 2, 3) if weight.dim() 4 else (1,)) num_keep int(weight.size(0) * (1 - prune_ratio)) _, keep_idx torch.topk(importance, max(1, num_keep)) keep_idx keep_idx.sort()[0] module.weight.data weight[keep_idx] if module.bias is not None: module.bias.data module.bias.data[keep_idx] # 应用量化此处为占位——实际需使用量化API quant_bits config.get(quant_bits, 32) if quant_bits 32: # 调用对应的量化函数 pass stats[compressed_params] sum( p.numel() for p in compressed_model.parameters() ) stats[compression_ratio] ( stats[original_params] / stats[compressed_params] ) return compressed_model, stats三、进化搜索算法的实现进化算法Evolutionary Algorithm是压缩NAS中最常用的搜索策略之一。其优势在于(1) 自然支持多目标优化精度延迟内存(2) 不需要可微的代理模型(3) 搜索结果具有多样性Pareto前沿。import random from typing import List, Callable from dataclasses import dataclass dataclass class Individual: 进化算法中的个体——一个压缩方案。 architecture: Dict # 压缩配置 accuracy: float 0.0 # 验证精度 model_size_mb: float 0.0 # 模型大小 latency_ms: float 0.0 # 推理延迟 fitness: float 0.0 # 综合适应度 class CompressionNAS_Evolution: 基于进化算法的压缩NAS搜索。 搜索流程 1. 初始化种群随机采样N个压缩方案 2. 评估每个个体训练toy model或使用代理模型 3. 选择锦标赛选择 → 高适应度个体 4. 交叉两个父代方案的层配置交换 5. 变异随机扰动部分层的配置 6. 重复2-5直至收敛 def __init__( self, search_space: CompressionSearchSpace, population_size: int 50, num_generations: int 20, mutation_rate: float 0.1, crossover_rate: float 0.5, # 多目标权重 accuracy_weight: float 1.0, size_weight: float 0.3, latency_weight: float 0.2 ): self.search_space search_space self.population_size population_size self.num_generations num_generations self.mutation_rate mutation_rate self.crossover_rate crossover_rate self.accuracy_weight accuracy_weight self.size_weight size_weight self.latency_weight latency_weight self.population: List[Individual] [] def initialize_population(self): 随机初始化种群。 self.population [] for _ in range(self.population_size): arch self.search_space.sample_random_architecture() self.population.append(Individual(architecturearch)) def compute_fitness(self, individual: Individual): 计算个体的综合适应度。 适应度 accuracy_weight * accuracy - size_weight * model_size_mb - latency_weight * latency_ms 目标: 最大化精度最小化体积和延迟。 individual.fitness ( self.accuracy_weight * individual.accuracy - self.size_weight * individual.model_size_mb - self.latency_weight * individual.latency_ms ) def tournament_select(self, tournament_size: int 3) - Individual: 锦标赛选择随机选k个个体返回最优者。 相比轮盘赌选择锦标赛选择对适应度scale不敏感 且更容易并行化。 candidates random.sample(self.population, tournament_size) return max(candidates, keylambda ind: ind.fitness) def crossover( self, parent1: Individual, parent2: Individual ) - Individual: 单点交叉随机选择一个层作为交叉点。 交叉点之前的配置来自parent1 交叉点之后的配置来自parent2。 layers self.search_space.compressible_layers crossover_point random.randint(1, len(layers) - 1) child_arch {} for i, layer in enumerate(layers): if i crossover_point: child_arch[layer] parent1.architecture[layer].copy() else: child_arch[layer] parent2.architecture[layer].copy() return Individual(architecturechild_arch) def mutate(self, individual: Individual) - Individual: 随机变异部分层的压缩配置。 每个层的每个动作有mutation_rate的概率 被替换为不同的随机选项。 mutated_arch {} for layer, config in individual.architecture.items(): mutated_config config.copy() for action, current_value in config.items(): if random.random() self.mutation_rate: # 随机选择不同的选项 options self.search_space.action_space[layer][action] new_value random.choice( [o for o in options if o ! current_value] ) mutated_config[action] new_value mutated_arch[layer] mutated_config return Individual(architecturemutated_arch) def search( self, evaluate_fn: Callable[[Dict], Tuple[float, float, float]] ) - List[Individual]: 执行完整的进化搜索。 Args: evaluate_fn: 评估函数 (architecture) → (accuracy, size_mb, latency) Returns: Pareto前沿个体列表 self.initialize_population() for generation in range(self.num_generations): # 评估当前种群 for individual in self.population: acc, size, lat evaluate_fn(individual.architecture) individual.accuracy acc individual.model_size_mb size individual.latency_ms lat self.compute_fitness(individual) # 精英保留最优个体直接进入下一代 self.population.sort(keylambda x: x.fitness, reverseTrue) elite_count max(2, self.population_size // 10) new_population self.population[:elite_count] # 生成下一代 while len(new_population) self.population_size: parent1 self.tournament_select() parent2 self.tournament_select() if random.random() self.crossover_rate: child self.crossover(parent1, parent2) else: child Individual( architectureparent1.architecture.copy() ) child self.mutate(child) new_population.append(child) self.population new_population # 评估最终种群并返回Pareto前沿 for individual in self.population: acc, size, lat evaluate_fn(individual.architecture) individual.accuracy acc individual.model_size_mb size individual.latency_ms lat return self._get_pareto_front() def _get_pareto_front(self) - List[Individual]: 提取Pareto前沿——多目标下无可改进的个体。 front [] for ind in self.population: dominated False for other in self.population: if (other.accuracy ind.accuracy and other.model_size_mb ind.model_size_mb and other.latency_ms ind.latency_ms and (other.accuracy ind.accuracy or other.model_size_mb ind.model_size_mb or other.latency_ms ind.latency_ms)): dominated True break if not dominated: front.append(ind) return front四、代理模型与加速搜索完整训练每个候选压缩方案在计算上是不可行的进化算法20代×50个体1000次训练。代理模型Surrogate Model通过在少量真实训练的基础上学习一个精度预测器来估计未训练的候选方案的精度。常用代理模型基于GCN图卷积网络的精度预测器——将模型架构编码为图节点算子边数据流通过图神经网络预测该架构在目标数据集上的精度。训练100-200个随机架构的真实精度后代理模型即可对新架构做出精度预测MAE约0.5-1.5%。五、总结压缩NAS将模型压缩从手工试错升级为自动化搜索核心价值在于当部署约束复杂多目标精度延迟内存能耗时机器比人更擅长在高维离散空间中寻找帕累托最优解。实践建议(1) 搜索空间不宜过大——合理限制每层的候选选项4-6个过多的选项会稀释搜索效率(2) 进化算法是压缩NAS的最佳起点——简单、鲁棒、天然支持多目标(3) 代理模型加速在搜索空间超过10^20时从可选变为必需(4) 搜索结果的最终精度仍需通过完整微调验证——代理模型预测的0.5% MAE可能掩盖真实的精度陷阱。