使用nni尋找最佳超參數

2023-08-27

前言

這篇筆記整理如何使用 NNI（Neural Network Intelligence）做超參數搜尋。範例任務是用 CNN 訓練 CIFAR-10，並透過 NNI 搜尋 learning rate、momentum 與 batch size。

作業檔案簡介

程式碼放在 ThreeMonth03/hyperparameter_tuning。

主要目錄如下：

config/：放 requirement.txt。
src/：放 source code，包含 cnn.py 與 nni_search.py。
log/：放 NNI experiment log，可以回放歷史 training 紀錄。

如何從頭復現 NNI Training

1
2
3

git clone https://github.com/ThreeMonth03/hyperparameter_tuning.git
cd hyperparameter_tuning
docker-compose up

接著在瀏覽器打開：

1	http://localhost:[your_port]

這裡建議不要使用 docker-compose up -d，否則 experiment log 可能不會被正常保存。實際部署時，也記得依照環境調整 port、container name 與 image name。

如何直接看 Training Log

如果只想查看既有 log，可以改用 nni_search.py 裡的 experiment.view：

1	experiment.view(experiment_id, port=8323, non_blocking=False)

操作流程：

Clone repo。
依照 nni_search.py 內的註解，關閉 training 設定，打開 experiment.view(...)。
執行 docker-compose up。
到 localhost:[your_port] 查看結果。

實驗設定

Hyperparameter	Search Space
`lr`	`0.0001 ~ 0.1`，log uniform
`momentum`	`0 ~ 1`，uniform
`batch_size`	`4`、`8`、`16`
Tuner	TPE

Result

Best hyperparameter：

lr: 0.0024724673142795927
momentum: 0.31344560117709097
batch_size: 8

Test Accuracy：65%

筆記

如何用 Python API 調 Hyperparameter

NNI 可以透過 terminal 指令或 Python API 控制 hyperparameter。以下是透過 Python API 設定 search space 與 experiment 的範例。

# nni_search.py
search_space = {
    'lr': {'_type': 'loguniform', '_value': [0.0001, 0.1]},
    'momentum': {'_type': 'uniform', '_value': [0, 1]},
    'batch_size': {"_type": "choice", "_value": [4, 8, 16]},
}

import nni
from nni.experiment import Experiment

experiment = Experiment('local')
experiment.config.trial_command = 'python src/cnn.py'
experiment.config.trial_code_directory = '.'
experiment.config.search_space = search_space
experiment.config.tuner.name = 'TPE'
experiment.config.tuner.class_args['optimize_mode'] = 'maximize'
experiment.config.max_trial_number = 50
experiment.config.trial_concurrency = 10
experiment.config.trial_gpu_number = 3
experiment.config.debug = True
experiment.config.experiment_working_directory = './log'
experiment.config.training_service.use_active_gpu = True
experiment.config.training_service.max_trial_number_per_gpu = 10

experiment.run(8323)
print(experiment.get_status())
print(experiment.get_job_statistics())
print(experiment.list_trial_jobs())

input('Press enter to quit')
experiment.stop()

被控制的 model 也要加入 NNI 參數讀取與回報結果的邏輯。

# cnn.py
import nni
#......
params = {
    'lr': 0.001,
    'momentum': 0,
    'batch_size': 4,
}
optimized_params = nni.get_next_parameter()
params.update(optimized_params)
print(params)
#......
epoches = 20
batch_size = params['batch_size']
lr = params['lr']
momentum = params['momentum']
#......
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')
nni.report_final_result(100 * correct // total)

小結

NNI 的好處是可以把「手動反覆調參」變成可重現的實驗流程。只要把 search space、tuner 與 training script 接好，就能自動化比較不同超參數組合，並保留 experiment log 供後續分析。

前言