proxylessnas
直接在目标硬件上搜索最优神经网络架构,消除代理任务与实际部署的性能鸿沟
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
直接在目标硬件上搜索最优神经网络架构,消除代理任务与实际部署的性能鸿沟
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
想象一个场景:你要为一款超轻量级手机开发图像识别模型,但不知道该选什么网络结构——ResNet 太重,MobileNet 勉强能跑但精度不够。传统做法是先设计再测量,或者用代理任务(比如小数据集)模拟。但 MIT 的研究者在 2019 年提出了一个更直接的问题:为什么不能直接在目标硬件和真实任务上,让算法自己决定网络长什么样? 这就是 ProxylessNAS 诞生的起点。
神经架构搜索(Neural Architecture Search, NAS)旨在让算法自动设计神经网络结构,取代人工调参。2018 年前的 NAS 方法普遍存在一个根本性问题——代理损失(Proxy Loss):
2019 年 ICLR 会议上,MIT Han Lab 的 Han Cai、Ligeng Zhu、Song Han 三位研究者发表了 ProxylessNAS,核心贡献是提出 Path-level Gradient-based NAS:

图1:ProxylessNAS 搜索框架。 算法在离散的操作空间中通过梯度下降直接优化,在每个决策点上用二值化门控选择最优路径,实现权重参数和架构参数的联合学习。
ProxylessNAS 的核心创新在于将 NAS 搜索空间从离散的节点空间转换为离散的路径空间,并利用 BinaryConnect 技术实现端到端梯度优化。
搜索空间包含以下几个维度的选择:
| 搜索维度 | 候选选项 |
|---|---|
| 卷积核大小 | [3, 5, 7] |
| 扩展比(Expansion Ratio) | [3, 6] |
| 路径是否跳过(Skip) | 是 / 否 |
| 算子类型 | Depthwise Conv, Mobile Inverted Bottleneck, 等 |
每个 Mobile Inverted Block 内的搜索空间极其庞大——论文中每个 block 有 10+ 个候选操作。通过 BinaryConnect 二值化方法,每个操作前加一个伯努利采样的二值门(0 或 1),最终输出为所有操作的加权和:
output = sum(g_i * op_i(x) for i in range(num_ops))
其中 g_i ∈ {0, 1},通过 Gumbel Softmax 或straight-through estimator(STE)实现不可导二值变量的梯度传递。
ProxylessNAS 的一大亮点是直接对目标硬件延迟建模。论文中使用延迟查找表(Latency Lookup Table):
L(w, a) = CrossEntropyLoss(w, a) + λ * Latency(a)
其中 Latency(a) 通过在目标硬件上测量每个算子的实际延迟并查表得到。搜索时,梯度回传中会包含对延迟的感知,从而让算法倾向于选择又快又准的结构。
这就是为什么 ProxylessNAS 搜索出来的 Mobile 版本,在特定手机上能稳定比 MobileNetV2 快——因为它不是在模拟延迟,而是在真实测量延迟。
ProxylessNAS 的影响力远超论文本身,以下几点值得特别关注:
ProxylessNAS 已集成到三大 AutoML 平台:
torch.hub.load('mit-han-lab/ProxylessNAS', 'proxyless_cpu', pretrained=True)这意味着该工作已经从学术论文变成了工业级工具链的一部分。
| 竞赛 | 名次 | 备注 |
|---|---|---|
| Visual Wake Words Challenge (CVPR 2019) | 第1名 | TF-lite 赛道 |
| Low Power Image Recognition Challenge (CVPR 2019) | 第3名 | 分类赛道 |

图2:ProxylessNAS 与主流模型在 ImageNet 上的 Top-1 精度与 GPU 延迟对比。 ProxylessNAS (GPU) 比 MobileNetV2 精度高 3.1%,同时延迟降低 20%。
代码库结构清晰,分为三个核心目录:
proxyless_nas/ # 预训练模型定义(PyTorch)
├── layers.py # 基础算子层(ConvLayer, DepthConvLayer, MBConvLayer 等)
├── nas_modules.py # NAS 核心模块(ProxylessNASNets 主网络类)
└── model_zoo.py # 模型动物园(CPU/GPU/Mobile 多种预设)
search/ # 架构搜索代码
├── nas_manager.py # 搜索配置与训练循环管理
├── imagenet_arch_search.py # ImageNet 上的搜索入口
└── run_manager.py # 训练运行管理(数据加载、优化器调度)
training/ # 模型微调代码
└── main.py # 微调脚本
核心代码逻辑(proxyless_nas/nas_modules.py):
class ProxylessNASNets(BasicUnit):
def __init__(self, first_conv, blocks, feature_mix_layer, classifier):
self.first_conv = first_conv
self.blocks = nn.ModuleList(blocks) # 可搜索的 Mobile Inverted Blocks
self.feature_mix_layer = feature_mix_layer
self.global_avg_pooling = nn.AdaptiveAvgPool2d(1)
self.classifier = classifier
def forward(self, x):
x = self.first_conv(x)
for block in self.blocks:
x = block(x) # 每个 block 内部有路径选择
x = self.feature_mix_layer(x)
x = self.global_avg_pooling(x)
x = x.view(x.size(0), -1)
x = self.classifier(x)
return x

图3:Mobile 设置下 ProxylessNAS 与 MobileNetV2 的性能对比。 在不同延迟约束下,ProxylessNAS 始终优于 MobileNetV2,体现了硬件感知搜索的显著优势。
PyTorch Hub 提供了四种预设模型,Python 环境安装 PyTorch 后,两行代码即可使用:
import torch
model = torch.hub.load('mit-han-lab/ProxylessNAS', 'proxyless_cpu', pretrained=True)
四种预设及适用场景:
| 预设 | 参数量 | Top-1 | 适用场景 |
|---|---|---|---|
| proxyless_cpu | ~21M | 74.5% | 服务器 / 高性能桌面 |
| proxyless_gpu | ~21M | 75.1% | GPU 服务器(延迟敏感) |
| proxyless_mobile | ~4.2M | 73.4% | 移动端 / 边缘设备 |
| proxyless_mobile14 | ~3.4M | 72.0% | 超轻量场景(延迟小于5ms) |
如需在特定硬件或数据集上搜索专属架构:
# 1. 准备 ImageNet 数据集
# 2. 修改搜索配置(search/imagenet_arch_search.py)
# 3. 启动搜索(需要多卡 GPU 机器)
python search/imagenet_arch_search.py --target hardware=cpu
搜索过程约 200 GPU hours,ImageNet 全量数据。对于企业内部定制化场景(如特定工业相机的图像分类),这是值得的投入。
| 场景 | GPU | 显存 | CPU | 内存 | 存储 |
|---|---|---|---|---|---|
| 推理(预训练模型) | 可选 | 2GB+ | 多核 | 4GB+ | 1GB |
| 微调 | 推荐 | 8GB+ | 多核 | 16GB+ | 10GB |
| 从头搜索 | 必须 | 32GB+(V100 8卡) | 多核 | 64GB+ | 50GB |

图4:ProxylessNAS 与其他 NAS 方法的搜索成本与最终精度对比。 图中可见 ProxylessNAS 在保持 SOTA 精度的同时,大幅降低了搜索计算成本。
ProxylessNAS 虽然开创性地解决了代理损失问题,但仍存在一些局限性:
这些问题在后续工作(如 Once-for-All、BigNAS、EAT-NAS)中得到了不同程度的解决。
ProxylessNAS 所在的时间节点(2019 年)恰好是 NAS 领域从小众学术方向走向工业标配的关键时期。它的核心贡献在于:

图5:针对不同硬件平台搜索专业化架构。 同一算法针对 CPU、GPU、Mobile 三种平台搜索出的结构各有不同,体现了硬件感知搜索的核心价值——专平台专用。

图6:ProxylessNAS 为不同平台搜索出的专业化架构可视化。 可以清晰看到针对不同硬件约束,网络的深度、宽度、连接模式均有显著差异。
ProxylessNAS 代表了一种重要的研究范式:不是从理论出发设计网络,而是从数据出发,让算法在真实约束下自己找到最优解。这种量体裁衣的思路深刻影响了后续 EfficientNet、Once-for-All、BigNAS 等工作,也为零样本 AutoML、硬件-算法联合优化等领域奠定了方法论基础。
对于 AI 爱好者来说,ProxylessNAS 是理解 NAS 思想的一个绝佳入口——代码清晰、文档完善、PyTorch Hub 一键可用。对于 AI 开发者来说,它提供了一套可直接复用的预训练模型,且其搜索框架本身也可作为企业内部 AutoML 系统的参考实现。
项目来源:MIT Han Lab / ProxylessNAS (ICLR 2019) | arXiv: 1812.00332