SPOS.md 4.64 KB
Newer Older
Chi Song's avatar
Chi Song committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
# 单路径 One-Shot (SPOS)

## 介绍

[Single Path One-Shot Neural Architecture Search with Uniform Sampling](https://arxiv.org/abs/1904.00420) 中提出的 one-shot NAS 方法,通过构造简化的通过统一路径采样方法训练的超网络来解决 One-Shot 模型训练的问题。这样所有架构(及其权重)都得到了完全且平等的训练。 然后,采用进化算法无需任何微调即可有效的搜索出性能最佳的体系结构。

在 NNI 上的实现基于 [官方 Repo](https://github.com/megvii-model/SinglePathOneShot). 实现了一个训练超级网络的 Trainer,以及一个利用 NNI 框架能力来加速进化搜索阶段的进化 Tuner。 还展示了

## 示例

此示例是论文中的搜索空间,使用 flops 限制来执行统一的采样方法。

[示例代码](https://github.com/microsoft/nni/tree/master/examples/nas/spos)

### 必需组件

由于使用了 DALI 来加速 ImageNet 的数据读取,需要 NVIDIA DALI >= 0.16。 [安装指南](https://docs.nvidia.com/deeplearning/sdk/dali-developer-guide/docs/installation.html)

[这里](https://1drv.ms/u/s!Am_mmG2-KsrnajesvSdfsq_cN48?e=aHVppN) ( [Megvii](https://github.com/megvii-model) 维护) 下载 flops 查找表。 将 `op_flops_dict.pkl``checkpoint-150000.pth.tar` (如果不需要重新训练超网络) 放到 `data` 目录中。

准备标准格式的 ImageNet (参考[这里的脚本](https://gist.github.com/BIGBALLON/8a71d225eff18d88e469e6ea9b39cef4))。 将其链接到 `data/imagenet` 会更方便。

准备好后,应具有以下代码结构:

```
spos
├── architecture_final.json
├── blocks.py
├── config_search.yml
├── data
│   ├── imagenet
│   │   ├── train
│   │   └── val
│   └── op_flops_dict.pkl
├── dataloader.py
├── network.py
├── readme.md
├── scratch.py
├── supernet.py
├── tester.py
├── tuner.py
└── utils.py
```

### 步骤 1. 训练超网络

```
python supernet.py
```

会将检查点导出到 `checkpoints` 目录中,为下一步做准备。

注意:数据加载的官方 Repo [与通常的方法有所不同](https://github.com/megvii-model/SinglePathOneShot/issues/5),使用了 BGR 张量,以及 0 到 255 之间的值来与自己的深度学习框架对齐。 选项 `--spos-preprocessing` 会模拟原始的使用行为,并能使用预训练的检查点。

### 步骤 2. 进化搜索

单路径 One-Shot 利用进化算法来搜索最佳架构。 tester 负责通过训练图像的子集来测试采样的体系结构,重新计算所有批处理规范,并在完整的验证集上评估架构。

为了使 Tuner 识别 flops 限制并能计算 flops,在 `tuner.py` 中创建了新的 `EvolutionWithFlops` Tuner,其继承于 SDK 中的 tuner。

要为 NNI 框架准备好搜索空间,首先运行

```
nnictl ss_gen -t "python tester.py"
```

将生成 `nni_auto_gen_search_space.json` 文件,这是搜索空间的序列化形式。

默认情况下,它将使用前面下载的 `checkpoint-150000.pth.tar`。 如果要使用从自行训练的检查点,在 `config_search.yml` 中的命令上指定 `---checkpoint`

然后使用进化 Tuner 搜索。

```
nnictl create --config config_search.yml
```

从每个 Epoch 导出的最终架构可在 Tuner 工作目录下的 `checkpoints` 中找到,默认值为 `$HOME/nni/experiments/your_experiment_id/log`

### 步骤 3. 从头开始训练

```
python scratch.py
```

默认情况下,它将使用 `architecture_final.json`. 该体系结构由官方仓库提供(转换成了 NNI 格式)。 通过 `--fixed-arc` 选项,可使用任何结构(例如,步骤 2 中找到的结构)。

## 参考

### PyTorch

```eval_rst
..  autoclass:: nni.nas.pytorch.spos.SPOSEvolution
    :members:

..  autoclass:: nni.nas.pytorch.spos.SPOSSupernetTrainer
    :members:

..  autoclass:: nni.nas.pytorch.spos.SPOSSupernetTrainingMutator
    :members:
```

## 已知的局限

* 仅支持 Block 搜索。 尚不支持通道搜索。
* 仅提供 GPU 版本。

## 当前重现结果

重现中。 由于官方版本和原始论文之间的不同,我们将当前结果与官方 Repo(我们运行的结果)和论文进行了比较。

* 进化阶段几乎与官方 Repo 一致。 进化算法显示出了收敛趋势,在搜索结束时达到约 65% 的精度。 但此结果与论文不一致。 详情参考[此 issue](https://github.com/megvii-model/SinglePathOneShot/issues/6)
* 重新训练阶段未匹配。 我们的重新训练代码,使用了作者发布的架构,获得了 72.14% 的准确率,与官方发布的 73.61%,和原始论文中的 74.3% 有一定差距。