这篇是发表在ICLR 2020的《Once-for-all: train one network and specialize it for efficient deployment》,作者是MIT的韩松团队,第一作者是Han Cai。
为了解决在不同设备和限制条件下神经网络高效推理的问题,一般会使用Hardware aware的NAS。作者提出使用一个once-for-all的网络,能够解耦训练和搜索过程,通过一次训练一个超大网络,再从中搜索子网络来得到限制条件下最高的性能。
方法
大网络设计(搜索空间)
- 多个units,每个unit里面是多层,只有第一层是stride=2
- 每个unit
- elastic layers: 任意数量的layers {2,3,4}
- elastic width: 任意数量的channels: expansion ratio{3,4,6}
- elastic kernel size: 任意kernel size {3,5,7}
- 网络数目$((3 × 3)^2 + (3 × 3)^3 + (3 × 3)^4)^5 ≈ 2 × 10^19$
- 5个Unit,最外是5次方
- 每个unit里面可以选择2,3,4三种block数量(幂次),加起来
- 每个layer可以3种kernel size选择,3种width选择,所以是3x3
- 另外,resolution还有25种可能
大网络训练
- sample各种配置的网络联合训练
- Progressive Shrinking
- resolution: sample 128-224大小的输入,全部训练过程都有
- 首先用最大的layers width kernel size训练,慢慢把不同配置加入sample池子
- 子网络的loss使用knowledge distillation
- width:大网络挑选L1最大的,作为子网络的initial,慢慢加入sample池
- depth: 取前n层
- kernel size: 7x7的center乘以一个变换矩阵得到子矩阵
搜索
- sample 16K个网络,在10K个image上训练accuracy predictor
- neural network twins
- 3层,每层有400个hidden units
- 输入是one-hot vector based on kernel size, input size and expand ratio,如果是skipped就zero vector
- neural network twins
- evolutionary search
- 优化目标accuracy/latency
结果
- ImageNet
- OFA 76% acc, 230M MAC
- OFA finetune#25epochs 76.4% 230M MAC
- OFA finetune#75epochs 76.9% 230M
- 比Mnasnet的74% 317M MAC好
- 比MobileNet V3-Large的75.2% 219M MAC好
- 训练时间为1200GPU hours
- 在Samsung S7 Note8/10 Pixel1/2测试
- 在1080Ti CPU TX2 ZU9EG FPGA ZU3EG上测试
- FPGA经验
- memory access是expensive的,computation是cheap的
- 通过OPs/Byte来作为一个指标
总结
once-for-all的思路是weight sharing的一个扩展,亮点在于graph不是predefine的,而是online generate的,借用了一些pruning的技术。有一个sharing的趋势吧,应该能做更多东西出来。