Optuna で PyTorch のハイパーパラメータチューニングを試みる（テーブルデータ回帰）

数ヶ月前の話ですが、Optuna が PyTorch のエコシステムに加わったとのことで前から試したいと思っていました。

https://note.com/aixtech/n/n5082ea1e2f47

MNIST での Colab でのチュートリアルが上記サイトで紹介されていましたので、テーブルデータの回帰問題で、ちょっと試してみようと思います。

おしながき

・データを用意

・ネットワークの定義

・Optuna で最適化

データを用意

今回は、カリフォルニアの住宅価格データセットを sklearn から持ってきて使います。

import pandas as pd
from sklearn.datasets import fetch_california_housing

datasets = fetch_california_housing()

df = pd.DataFrame(datasets.data, columns=datasets.feature_names)

print(df.shape)

'''
(20640, 8)
'''

import pandas as pd

from sklearn.datasets import fetch_california_housing

datasets = fetch_california_housing()

df = pd.DataFrame(datasets.data, columns=datasets.feature_names)

print(df.shape)

'''

(20640, 8)

'''

全部で 20640 件、8つの説明変数を持つデータセットです。

データローダーの作成までがっとやってしまいます。

import torch
import torch.nn as nn  
import torch.nn.functional as F

torch.manual_seed(2020)

x = torch.tensor(datasets.data, dtype=torch.float32)
t = torch.tensor(datasets.target, dtype=torch.float32)

tensorDataset = torch.utils.data.TensorDataset(x, t)

n_train = int(len(tensorDataset) * 0.7)
n_test = len(tensorDataset) - n_train 

train, test = torch.utils.data.random_split(tensorDataset, [n_train, n_test])

batch_size = 1024

train_dataloader = torch.utils.data.DataLoader(train, batch_size, shuffle=True)
test_dataloader = torch.utils.data.DataLoader(test, batch_size)

import torch

import torch.nn as nn

import torch.nn.functional as F

torch.manual_seed(2020)

x = torch.tensor(datasets.data, dtype=torch.float32)

t = torch.tensor(datasets.target, dtype=torch.float32)

tensorDataset = torch.utils.data.TensorDataset(x, t)

n_train = int(len(tensorDataset) * 0.7)

n_test = len(tensorDataset) - n_train

train, test = torch.utils.data.random_split(tensorDataset, [n_train, n_test])

batch_size = 1024

train_dataloader = torch.utils.data.DataLoader(train, batch_size, shuffle=True)

test_dataloader = torch.utils.data.DataLoader(test, batch_size)

ネットワークの定義

以下の 4 つをチューニングする前提

最適化手法
学習率
中間層の数
中間層のノードの数

中間層の数を可変にするために、層を input_layer, hiddne_layer, outpu_layer の 3 つの塊にわけて、hidden_layer は for ループで層を任意の数追加できるようにする。

class Net(nn.Module):

  def __init__(self, input_dim, hidden_dim, n_layers):
    super().__init__()
    self.input_dim = input_dim
    self.hidden_dim = hidden_dim
    self.n_layers = n_layers

    self.input_layer = nn.Sequential(
        nn.BatchNorm1d(input_dim),
        nn.Linear(input_dim, hidden_dim),
        nn.ReLU(inplace=True)
    )

    middle = []
    for _ in range(n_layers):
      middle.append(nn.BatchNorm1d(hidden_dim),)
      middle.append(nn.Linear(hidden_dim, hidden_dim))
      middle.append(nn.ReLU(inplace=True))

    self.middle_layers = nn.Sequential(*middle)

    self.output_layer = nn.Linear(hidden_dim, 1)


  def forward(self, x):
    x = self.input_layer(x)
    x = self.middle_layers(x)
    x = self.output_layer(x)
    return x

class Net(nn.Module):

def __init__(self, input_dim, hidden_dim, n_layers):

super().__init__()

self.input_dim = input_dim

self.hidden_dim = hidden_dim

self.n_layers = n_layers

self.input_layer = nn.Sequential(

nn.BatchNorm1d(input_dim),

nn.Linear(input_dim, hidden_dim),

nn.ReLU(inplace=True)

)

middle = []

for _ in range(n_layers):

middle.append(nn.BatchNorm1d(hidden_dim),)

middle.append(nn.Linear(hidden_dim, hidden_dim))

middle.append(nn.ReLU(inplace=True))

self.middle_layers = nn.Sequential(*middle)

self.output_layer = nn.Linear(hidden_dim, 1)

def forward(self, x):

x = self.input_layer(x)

x = self.middle_layers(x)

x = self.output_layer(x)

return x

Optuna で最適化

optuna では、objective 関数内で学習ループを書き、返り値に与える関数を最小化する。

今回は、validation loss を最小化するようにチューニングしてもらいます。

調整したいハイパーパラメータがどんな値を取るかによって suggestほにゃららメソッドは変わる。

max_epoch = 30
criterion = nn.MSELoss()

def objective(trial):
  device = "cuda" if torch.cuda.is_available() else "cpu"

  n_layers = trial.suggest_int("n_layers", 1, 5)
  hidden_dim = trial.suggest_int("hidden_dim", 8, 128, 8)
  optimizer = trial.suggest_categorical("optimizer", ["Adam", "RMSprop", "SGD"])
  learning_rate = trial.suggest_discrete_uniform("learning_rate", 1e-5, 1e-1, q=1e-5)

  net = Net(input_dim=8, hidden_dim=hidden_dim, n_layers=n_layers).to(device)
  optimizer = getattr(torch.optim, optimizer)(net.parameters(), lr=learning_rate)

  # train 
  net.train()
  for epoch in range(max_epoch):
    for batch in train_dataloader:
      x, t = batch
      x.to(device)
      t.to(device)

      optimizer.zero_grad()
      y = net(x)
      loss = criterion(y, t)
      loss.backward()
      optimizer.step()
    
  # validation
  net.eval()
  validation_loss = 0.0
  with torch.no_grad():
    for batch in test_dataloader:
      x, t = batch
      x.to(device)
      t.to(device)
      y = net(x)
      loss = criterion(y, t)
      validation_loss += loss.item() * x.size(0)

    validation_loss = validation_loss / len(test_dataloader.dataset)

  return validation_loss

max_epoch = 30

criterion = nn.MSELoss()

def objective(trial):

device = "cuda" if torch.cuda.is_available() else "cpu"

n_layers = trial.suggest_int("n_layers", 1, 5)

hidden_dim = trial.suggest_int("hidden_dim", 8, 128, 8)

optimizer = trial.suggest_categorical("optimizer", ["Adam", "RMSprop", "SGD"])

learning_rate = trial.suggest_discrete_uniform("learning_rate", 1e-5, 1e-1, q=1e-5)

net = Net(input_dim=8, hidden_dim=hidden_dim, n_layers=n_layers).to(device)

optimizer = getattr(torch.optim, optimizer)(net.parameters(), lr=learning_rate)

# train

net.train()

for epoch in range(max_epoch):

for batch in train_dataloader:

x, t = batch

x.to(device)

t.to(device)

optimizer.zero_grad()

y = net(x)

loss = criterion(y, t)

loss.backward()

optimizer.step()

# validation

net.eval()

validation_loss = 0.0

with torch.no_grad():

for batch in test_dataloader:

x, t = batch

x.to(device)

t.to(device)

y = net(x)

loss = criterion(y, t)

validation_loss += loss.item() * x.size(0)

validation_loss = validation_loss / len(test_dataloader.dataset)

return validation_loss

import optuna

study = optuna.create_study()
study.optimize(objective, n_trials=50)

import optuna

study = optuna.create_study()

study.optimize(objective, n_trials=50)

結果は、以下のようにして確認することができる。

# 一番良い結果
study.best_trial

'''
{'hidden_dim': 72,
 'learning_rate': 0.09883,
 'n_layers': 4,
 'optimizer': 'Adam'}
'''

# データフレームで確認
study.trials_dataframe()

# values が objective 関数の返り値
study.trials_dataframe().sort_values("value").head(10)