如何使用Optuna优化AI大模型训练？自动化调参的详细教程

objective

trial.suggest_xxx()

trial.report(metric_value, step)

optuna.create_study()

study.optimize(objective, n_trials=...)

study.best_trial.params

study.best_trial.value

trial.suggest_float("lr", 1e-5, 1e-2, log=True)

trial.suggest_int("batch_size", 32, 256, step=32)

trial.suggest_categorical()

import optuna
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# 简化的大模型模拟：一个多层感知机
class SimpleMLP(nn.Module):
    def __init__(self, input_dim, num_classes, n_layers, hidden_dim, dropout_rate):
        super(SimpleMLP, self).__init__()
        layers = []
        layers.append(nn.Linear(input_dim, hidden_dim))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(dropout_rate))
        for _ in range(n_layers - 1):
            layers.append(nn.Linear(hidden_dim, hidden_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout_rate))
        layers.append(nn.Linear(hidden_dim, num_classes))
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

def objective(trial):
    # 1. 超参数采样
    lr = trial.suggest_float("lr", 1e-5, 1e-2, log=True)
    optimizer_name = trial.suggest_categorical("optimizer", ["Adam", "SGD", "RMSprop"])
    n_layers = trial.suggest_int("n_layers", 2, 5) # 模型层数
    hidden_dim = trial.suggest_int("hidden_dim", 64, 512, step=64) # 隐藏层维度
    dropout_rate = trial.suggest_float("dropout_rate", 0.1, 0.5)
    batch_size = trial.suggest_int("batch_size", 64, 256, step=64)
    epochs = trial.suggest_int("epochs", 5, 20) # 模拟训练epoch，实际大模型可能更多

    # 2. 数据准备 (简化，实际大模型会涉及更复杂的数据加载和预处理)
    X, y = make_classification(n_samples=1000, n_features=20, n_classes=2, random_state=42)
    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

    X_train_t = torch.tensor(X_train, dtype=torch.float32)
    y_train_t = torch.tensor(y_train, dtype=torch.long)
    X_val_t = torch.tensor(X_val, dtype=torch.float32)
    y_val_t = torch.tensor(y_val, dtype=torch.long)

    train_dataset = TensorDataset(X_train_t, y_train_t)
    val_dataset = TensorDataset(X_val_t, y_val_t)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

    # 3. 模型、优化器、损失函数初始化
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = SimpleMLP(X.shape[1], 2, n_layers, hidden_dim, dropout_rate).to(device)

    if optimizer_name == "Adam":
        optimizer = optim.Adam(model.parameters(), lr=lr)
    elif optimizer_name == "SGD":
        optimizer = optim.SGD(model.parameters(), lr=lr)
    else: # RMSprop
        optimizer = optim.RMSprop(model.parameters(), lr=lr)

    criterion = nn.CrossEntropyLoss()

    # 4. 训练循环与剪枝
    for epoch in range(epochs):
        model.train()
        for batch_X, batch_y in train_loader:
            batch_X, batch_y = batch_X.to(device), batch_y.to(device)
            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()

        # 验证模型
        model.eval()
        val_preds = []
        val_true = []
        with torch.no_grad():
            for batch_X, batch_y in val_loader:
                batch_X, batch_y = batch_X.to(device), batch_y.to(device)
                outputs = model(batch_X)
                _, predicted = torch.max(outputs.data, 1)
                val_preds.extend(predicted.cpu().numpy())
                val_true.extend(batch_y.cpu().numpy())

        val_accuracy = accuracy_score(val_true, val_preds)

        # 向Optuna报告中间结果，并检查是否需要剪枝
        trial.report(val_accuracy, epoch)
        if trial.should_prune():
            raise optuna.exceptions.TrialPruned()

    return val_accuracy # 返回最终的验证集准确率，Optuna将最大化此值

trial.report(current_metric_value, current_step)

if trial.should_prune(): raise optuna.exceptions.TrialPruned()

trial.report