Ray tune 를 사용해서 pytorch hyperparameter 최적화하기 (feat. diabetes)

Ray란?

홈페이지에는 아래와 같이 소개하고 있지만, 결국 AI 관련된 Python workload 를 좀더 간편하게 개발 할수 있도록 도와주는 프레임워크이다.

Effortlessly scale your most complex workloads
Ray is an open-source unified compute framework that makes it easy to scale AI and Python workloads — from reinforcement learning to deep learning to tuning, and model serving. Learn more about Ray’s rich set of libraries and integrations.

대표 기능을 살펴보면 아래와 같다. Deep Learning 자체에 대한 기술이라기 보다는, Hyperparameter Tuning, Serving, Data Processing 등 효과적인 모델 개발을 도와주는 것을 목표로 한다.

ㅇ

Hyperparameter Tuning

이중 Hypterparameter Tuning 기능을 사용해보자

MLP Regression 모델 공통코드 만들기

일단 학습을 위한 공통 코드를 만든다. 최적화할 hyperparameter 를 미리 정의하고 해당 값이 변경되어도 정상적인 학습을 할수 있도록  model, optimizer, loss function 등을 정의해야 한다. 본 예시에서는 최적화할 hyperparameter를 3가지로 정했다. (learning rate, layer2 in, layer2_out)

• 데이터 준비 (sklearn의 diabetes 데이터)
• train: 0.8, test: 0.2

from sklearn import datasets
from sklearn.model_selection import train_test_split

from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader

data = datasets.load_diabetes()
x_train, x_test, y_train, y_test = train_test_split(data.data, data.target, test_size=0.2, random_state=0)

train_set = TensorDataset(torch.tensor(x_train).float(), torch.tensor(y_train).long())
train_loader = DataLoader(train_set, batch_size=128)

x_test = torch.tensor(x_test).float()
y_test = torch.tensor(y_test).float()

• model, optimizer, loss function 정의
• 각 객체 생성시점에 hyperparameter 를 사용할수 있도록 정의
  • optimizer 의 lr
  • model 의 l2_in, 2_out

import torch
from torch import nn

def get_optimizer(model, lr):
  return torch.optim.SGD(model.parameters(), lr=lr)

def get_loss_fn():
  return torch.nn.MSELoss()

class MyModel(nn.Module):
  def __init__(self, l2_in, l2_out):
    super().__init__()

    self._model = nn.Sequential(
        nn.Linear(x.size(-1), l2_in),
        nn.ReLU(),
        nn.Linear(l2_in, l2_out),
        nn.ReLU(),
        nn.Linear(l2_out, 1)
    )

  def forward(self, x):
    pred = self._model(x)
    return pred

• train 함수 정의 (이전글 참고)

import torch
from torch import nn
  
def train(loader, model, optimizer, loss_fn):
  epoch=1000
  print_num=100

  for e in range(epoch):
    for i, (x, y) in enumerate(loader):
      pred = model(x)
      loss = loss_fn(pred, y)

      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

      if e % print_num == 0:
        print(e, i, loss.item())
  print(e, i, loss.item())

Ray로 Hyperparamer Tuning 하기

해당 기능을 사용하기 위해서는 아래와 같은 작업이 필요하다.

1. Ray.tune.Tuner 객체 생성
2. tunner.fit() 실행

Ray.tune.Tuner 객체 생성

Ray.tune.Tuner 객체를 생성하기 위해서는 아래와 같이 4가지를 정의해야 한다.

• trainable – The trainable to be tuned.
• param_space – Search space of the tuning job. One thing to note is that both preprocessor and dataset can be tuned here.
• tune_config – Tuning algorithm specific configs. Refer to ray.tune.tune_config.TuneConfig for more info.
• run_config – Runtime configuration that is specific to individual trials. If passed, this will overwrite the run config passed to the Trainer, if applicable. Refer to ray.air.config.RunConfig for more info.

먼저 trainable 을 살펴보자. 실제 학습을 수행하는 함수이다. config 를 파라미터로 받으며, 이는 Tuner 객체 생성시 param_space 로 전달된 값이다. 아래와 같이 정의한다.

• config 의 hyperparameter 를 trainable 에서 적절하게 사용한다.
• 마지막 최종 결과값을 tune.report()을 통해 tune 에게 전달한다.
• reporting 되는 값은 test 데이터셋을 사용하여 계산하며, tensor 가 아닌 scalar 값으로 변환한다.
• while True 로 인해 마치 무한 loop 처럼 보이지만, 아래 tune_config.num_samples 설정으로 정상적으로 종료된다.

def trainable_regression(config):
    model = MyModel(config['l2_in'], config['l2_out'])
    optimizer = get_optimizer(model, lr=config['lr'])
    loss_fn = get_loss_fn()

    while True:
        train(train_loader, model, optimizer, loss_fn)
        loss_val = loss_fn(model(x_test), y_test)
        tune.report(loss=float(loss_val))

Tuner 객체 생성을 위한 그 외의 파라미터들은 아래와 같이 정의한다.

• param_space(param_space): 최적화할 hyperparameter 를 정의한다. parmaeter 성격에 따라 적절한 type 으로 정의한다.
• tune_config(tune_config): tuning 동작을 위한 config를 정의한다.
• run_config(run_config): trainable 실행을 위한 config를 정의한다.
• trainable_regression(trainable): 실제 학습을 실행할 함수이다.

from ray import tune, air
from ray.air import session

param_space = {
    "l2_in": tune.grid_search([10, 100, 200]),
    "l2_out": tune.grid_search([10, 100, 200]),
    "lr": tune.uniform(1e-10, 1e-7)
    }

tune_config=tune.TuneConfig(
    metric="loss",
    mode="min",
    num_samples=3
)

run_config=air.RunConfig(
    stop={"training_iteration": 2}
)

tuner = tune.Tuner(
    train_regression,
    tune_config=tune_config,
    run_config=run_config,
    param_space=param_space,
)

hyperparameter 최적화

• tuner.fit() 함수를 사용한다.

results = tuner.fit()
print("Best config is:", results.get_best_result().config)

최적화 결과

• 최종 결과는 아래와 같다
• Best config is: {'l2_in': 10, 'l2_out': 200, 'lr': 8.027810161479782e-08}
• loss 가 5838~29119으로 hyperparameter에 따라 최대 5배 정도의 차이가 발생함을 확인할수 있다.
• 3개 parameter에 대해 각 3가지 경우만 적용해도 3*3*3=27번의 학습이 발생한다.

Current time: 2023-06-14 05:09:14 (running for 00:41:41.74)
Using FIFO scheduling algorithm.
Logical resource usage: 1.0/2 CPUs, 0/1 GPUs
Current best trial: c81fa_00024 with loss=5838.0498046875 and parameters={'l2_in': 10, 'l2_out': 200, 'lr': 8.027810161479782e-08}
Result logdir: /root/ray_results/train_regression_2023-06-14_04-27-22
Number of trials: 27/27 (27 TERMINATED)
+------------------------------+------------+-------------------+---------+----------+-------------+--------+------------------+----------+
| Trial name                   | status     | loc               |   l2_in |   l2_out |          lr |   iter |   total time (s) |     loss |
|------------------------------+------------+-------------------+---------+----------+-------------+--------+------------------+----------|
| train_regression_c81fa_00000 | TERMINATED | 172.28.0.12:17553 |      10 |       10 | 7.24429e-08 |      2 |          187.108 | 28353.1  |
| train_regression_c81fa_00001 | TERMINATED | 172.28.0.12:17560 |     100 |       10 | 2.21732e-08 |      2 |          193.757 | 28967.5  |
| train_regression_c81fa_00002 | TERMINATED | 172.28.0.12:17553 |     200 |       10 | 5.09034e-08 |      2 |          186.721 | 27868.7  |
| train_regression_c81fa_00003 | TERMINATED | 172.28.0.12:17560 |      10 |      100 | 2.14678e-09 |      2 |          182.668 | 29075.7  |
| train_regression_c81fa_00004 | TERMINATED | 172.28.0.12:17553 |     100 |      100 | 1.9599e-09  |      2 |          195.521 | 29030.3  |
| train_regression_c81fa_00005 | TERMINATED | 172.28.0.12:17560 |     200 |      100 | 2.40337e-08 |      2 |          211.889 | 28602.5  |
| train_regression_c81fa_00006 | TERMINATED | 172.28.0.12:17553 |      10 |      200 | 6.80341e-08 |      2 |          182.721 |  6092.47 |
| train_regression_c81fa_00007 | TERMINATED | 172.28.0.12:17560 |     100 |      200 | 6.46715e-08 |      2 |          217.872 | 23455    |
| train_regression_c81fa_00008 | TERMINATED | 172.28.0.12:17553 |     200 |      200 | 5.15208e-08 |      2 |          237.965 | 20935.9  |
| train_regression_c81fa_00009 | TERMINATED | 172.28.0.12:17560 |      10 |       10 | 6.04327e-08 |      2 |          167.508 | 28537.4  |
| train_regression_c81fa_00010 | TERMINATED | 172.28.0.12:17560 |     100 |       10 | 9.2759e-08  |      2 |          155.346 |  6058.17 |
| train_regression_c81fa_00011 | TERMINATED | 172.28.0.12:17553 |     200 |       10 | 2.23245e-09 |      2 |          157.803 | 29119.5  |
| train_regression_c81fa_00012 | TERMINATED | 172.28.0.12:17560 |      10 |      100 | 7.45451e-08 |      2 |          162.223 | 25000    |
| train_regression_c81fa_00013 | TERMINATED | 172.28.0.12:17553 |     100 |      100 | 3.38622e-08 |      2 |          176.118 | 28528.1  |
| train_regression_c81fa_00014 | TERMINATED | 172.28.0.12:17560 |     200 |      100 | 8.11087e-08 |      2 |          188.447 |  5849.64 |
| train_regression_c81fa_00015 | TERMINATED | 172.28.0.12:17553 |      10 |      200 | 6.11654e-08 |      2 |          158.374 | 25694.5  |
| train_regression_c81fa_00016 | TERMINATED | 172.28.0.12:17560 |     100 |      200 | 2.79419e-08 |      2 |          188.947 | 28518.6  |
| train_regression_c81fa_00017 | TERMINATED | 172.28.0.12:17553 |     200 |      200 | 2.68204e-08 |      2 |          215.191 | 28518.8  |
| train_regression_c81fa_00018 | TERMINATED | 172.28.0.12:17560 |      10 |       10 | 9.99507e-08 |      2 |          149.954 | 26321.8  |
| train_regression_c81fa_00019 | TERMINATED | 172.28.0.12:17553 |     100 |       10 | 8.069e-08   |      2 |          152.7   | 27506.7  |
| train_regression_c81fa_00020 | TERMINATED | 172.28.0.12:17560 |     200 |       10 | 4.84234e-08 |      2 |          164.779 | 26451    |
| train_regression_c81fa_00021 | TERMINATED | 172.28.0.12:17553 |      10 |      100 | 2.5675e-08  |      2 |          155.723 | 28607    |
| train_regression_c81fa_00022 | TERMINATED | 172.28.0.12:17560 |     100 |      100 | 1.82331e-08 |      2 |          172.894 | 28805.3  |
| train_regression_c81fa_00023 | TERMINATED | 172.28.0.12:17553 |     200 |      100 | 6.66978e-08 |      2 |          187.397 | 11189.8  |
| train_regression_c81fa_00024 | TERMINATED | 172.28.0.12:17560 |      10 |      200 | 8.02781e-08 |      2 |          156.361 |  5838.05 |
| train_regression_c81fa_00025 | TERMINATED | 172.28.0.12:17553 |     100 |      200 | 6.94886e-08 |      2 |          187.479 | 11939.3  |
| train_regression_c81fa_00026 | TERMINATED | 172.28.0.12:17560 |     200 |      200 | 2.47206e-08 |      2 |          156.352 | 28505.4  |
+------------------------------+------------+-------------------+---------+----------+-------------+--------+------------------+----------+


Best config is: {'l2_in': 10, 'l2_out': 200, 'lr': 8.027810161479782e-08}