複数のGPUでOptunaを走らせる

せっかく複数のGPUがあるんだから、シングルGPUで動く深層学習のハイパラチューニングを並列にできないか試してみた。

ダメな例

以下、ChatGPTに生成させたコードを若干手直ししたもの。これでも小さなモデルだと動いてしまうが、複数のtrialが同じGPUに乗ってしまう可能性があるため、GPUのマシンがバグったりCUDA out of memoryのエラーが出ることがある。

"""
ChatGPTに吐かせたコード、ダメな例
"""

import optuna
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import multiprocessing

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

def objective(trial):
    
    # 例としてランダムデータを使用
    X_train = torch.randn(1000, 20)
    y_train = torch.randint(0, 2, (1000,))

    train_dataset = TensorDataset(X_train, y_train)
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    
    # ハイパーパラメータのサンプリング
    hidden_dim = trial.suggest_int('hidden_dim', 16, 128)
    lr = trial.suggest_loguniform('lr', 1e-5, 1e-1)
    
    # モデルの初期化とGPU設定
    # ここがNG!!!
    device = torch.device(f'cuda:{trial.number % torch.cuda.device_count()}')
    model = SimpleNet(input_dim=20, hidden_dim=hidden_dim, output_dim=1).to(device)
    
    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    
    # 訓練ループ
    for epoch in range(10):  # 例として10エポック
        model.train()
        for batch in train_loader:
            X_batch, y_batch = batch
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            
            optimizer.zero_grad()
            outputs = model(X_batch)
            loss = criterion(outputs.squeeze(), y_batch.float())
            loss.backward()
            optimizer.step()

    # 評価 (例として訓練損失を返す)
    return loss.item()

def run_optimization(storage_url, study_name):
    study = optuna.create_study(study_name=study_name, storage=storage_url, direction='minimize')
    study.optimize(objective, n_trials=100)

if __name__ == '__main__':
    storage_url = 'sqlite:///example.db'  # SQLiteの例
    study_name = 'distributed-study'
    
    # プロセス数（GPUの数と同じに設定）
    n_procs = torch.cuda.device_count()
    
    processes = []
    for _ in range(n_procs):
        p = multiprocessing.Process(target=run_optimization, args=(storage_url, study_name))
        p.start()
        processes.append(p)
    
    for p in processes:
        p.join()

良い例

OptunaのGithubにあったこのissueを参考に書き直すと、目的関数が利用したGPUの整理までやってくれるので、重複することがなくなる。

# 上記に加え追加でインポートしたモジュール
from multiprocessing import Manager
from joblib import parallel_backend
from datetime import datetime, UTC, timedelta
import logging

class Objective:
    def __init__(self, gpu_queue):
        self.gpu_queue = gpu_queue
    
    def __call__(self, trial):
                
        # モデルの初期化とGPU設定
        # 利用するGPUの取得
        gpu_id = self.gpu_queue.get()
        device = torch.device(f'cuda:{gpu_id}')
        print(f"Trial {trial.number}, Using Decice cuda:{gpu_id}")
        model = SimpleNet(input_dim=20, hidden_dim=hidden_dim, output_dim=1).to(device)
        
        ###省略：上記と同様にデータセットの構築や最適化関数の設定を行い学習させる###

        # GPUが使い終わったと教える
        self.gpu_queue.put(gpu_id)

        # 評価 (例として訓練損失を返す)
        return loss.item()

if __name__ == '__main__':
    storage_url = 'sqlite:///example.db'  # SQLiteの例

    # できれば一意に定まる名前の方がいい
    utc_time = datetime.now(UTC)
    execution_time_str = (utc_time + timedelta(hours=9)).strftime('%Y%m%d_%H%M')
    study_name = f'tutorial_{execution_time_str}'
    
    study = optuna.create_study(study_name=study_name, storage=storage_url, direction='minimize')
    
    # optunaのログを保持しておきたい場合はこのコードを書く
    optuna.logging.get_logger("optuna").addHandler(logging.FileHandler(f'optuna.log'))
    
    with Manager() as manager:
        gpu_queue = manager.Queue()
        for i in range(4,8):
            gpu_queue.put(i)
            
        with parallel_backend("multiprocessing", n_jobs=4):
            study.optimize(Objective(gpu_queue), n_trials=100, n_jobs=4)