PyTorch Lightning 使用 Optuna 进行超参数优化
作者 | Ideogram 提供图片
PyTorch Lightning 是近年来发布的一个高级替代方案，用于替代经典的 PyTorch 深度学习建模库。它简化了模型的训练、验证和部署过程。在进行超参数优化时，即寻找能最大化给定任务性能的最佳模型参数集，Optuna 可以是一个与 PyTorch Lightning 结合的强大工具，因为它易于集成，并提供了高效的搜索算法，可以在大量可能的配置中找到模型的最佳设置。
本文将介绍如何联合使用 PyTorch Lightning 和 Optuna 来指导深度学习模型的超参数优化过程。建议您具备实用的神经网络构建和训练基础知识，最好是使用 PyTorch。
分步流程
过程首先需要安装和导入一系列必要的库和模块，包括 PyTorch Lightning 和 Optuna。初始安装过程可能需要一些时间来完成。
|
1 2 3 |
pip install pytorch_lightning pip install optuna pip install optuna-integration[pytorch_lightning] |
接下来，进行大量导入
|
1 2 3 4 5 6 7 8 9 10 11 12 13 |
import os import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, random_split from torchvision import datasets, transforms import pytorch_lightning as pl from pytorch_lightning.callbacks import EarlyStopping from pytorch_lightning.loggers import TensorBoardLogger import optuna from optuna.integration import PyTorchLightningPruningCallback |
在使用 PyTorch Lightning 构建神经网络模型时，设置随机种子以保证结果的可复现性是一种常见做法。您可以在代码开头，导入语句之后，通过添加 pl.seed_everything(42) 来实现这一点。
接下来，我们通过创建一个继承自 pl.LightningModule 的类来定义神经网络模型架构：这是 Lightning 中与 PyTorch 的 Module 类相对应的部分。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
class MNISTClassifier(pl.LightningModule): def __init__(self, layer_1_size=128, layer_2_size=256, learning_rate=1e-3, dropout_rate=0.5): super().__init__() self.save_hyperparameters() # 神经网络架构 self.layer_1 = nn.Linear(28 * 28, self.hparams.layer_1_size) self.layer_2 = nn.Linear(self.hparams.layer_1_size, self.hparams.layer_2_size) self.layer_3 = nn.Linear(self.hparams.layer_2_size, 10) self.dropout = nn.Dropout(self.hparams.dropout_rate) def forward(self, x): # 展平层 batch_size, _, _, _ = x.size() x = x.view(batch_size, -1) # 前向传播 x = F.relu(self.layer_1(x)) x = self.dropout(x) x = F.relu(self.layer_2(x)) x = self.dropout(x) x = self.layer_3(x) return F.log_softmax(x, dim=1) def training_step(self, batch, batch_idx): x, y = batch logits = self(x) loss = F.nll_loss(logits, y) self.log('train_loss', loss, prog_bar=True) return loss def validation_step(self, batch, batch_idx): x, y = batch logits = self(x) loss = F.nll_loss(logits, y) preds = torch.argmax(logits, dim=1) acc = accuracy(preds, y) self.log('val_loss', loss, prog_bar=True) self.log('val_acc', acc, prog_bar=True) return {'val_loss': loss, 'val_acc': acc} def test_step(self, batch, batch_idx): x, y = batch logits = self(x) loss = F.nll_loss(logits, y) preds = torch.argmax(logits, dim=1) acc = accuracy(preds, y) self.log('test_loss', loss, prog_bar=True) self.log('test_acc', acc, prog_bar=True) return {'test_loss': loss, 'test_acc': acc} def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate) return optimizer |
我们将类命名为 MNISTClassifier，因为它是一个简单的前馈神经网络分类器，我们将使用 MNIST 数据集对其进行训练，以进行低分辨率图像分类。它包含少量线性层，前面有一个展平二维图像数据的输入层，中间是 ReLU 激活函数。该类还定义了用于前向传播的 forward() 方法，以及模拟训练、测试和验证步骤的其他方法。
在新建的类之外，下面的函数也将有助于计算一组多项预测的平均准确率
|
1 2 |
def accuracy(preds, y): return (preds == y).float().mean() |
以下函数建立了一个数据准备管道，我们稍后将将其应用于 MNIST 数据集。它将数据集转换为张量，根据该数据集的先验已知统计数据对其进行归一化，下载并将训练数据分割为训练集和验证集，并为三个数据子集中的每个子集创建一个 DataLoader 对象。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
def prepare_data(): transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) # MNIST 数据集的均值和标准差 ]) # 下载数据集 mnist_train = datasets.MNIST('data', train=True, download=True, transform=transform) mnist_test = datasets.MNIST('data', train=False, download=True, transform=transform) # 将原始训练数据分割为训练集和验证集 mnist_train, mnist_val = random_split(mnist_train, [55000, 5000]) # 为 PyTorch 数据管理创建 DataLoader train_loader = DataLoader(mnist_train, batch_size=64, shuffle=True) val_loader = DataLoader(mnist_val, batch_size=64) test_loader = DataLoader(mnist_test, batch_size=64) |
objective() 函数是 Optuna 为定义超参数优化框架提供的核心元素。我们通过首先定义搜索空间来定义它：即需要尝试的超参数及其可能的值。这可以包括与神经网络层相关的架构超参数，也可以包括与训练算法相关的超参数，例如学习率和用于对抗过拟合等问题的 dropout 率。
在此函数中，我们尝试为每个可能的超参数设置初始化和训练模型，并包含一个用于提前停止的回调，以防模型过早稳定。与标准的 PyTorch 一样，Trainer 用于模拟训练过程。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
# 定义 Optuna 的目标函数 def objective(trial): # 设置要优化的超参数 layer_1_size = trial.suggest_int('layer_1_size', 64, 256) layer_2_size = trial.suggest_int('layer_2_size', 128, 512) learning_rate = trial.suggest_float('learning_rate', 1e-4, 1e-2, log=True) dropout_rate = trial.suggest_float('dropout_rate', 0.2, 0.7) # 使用试验的超参数创建模型 model = MNISTClassifier( layer_1_size=layer_1_size, layer_2_size=layer_2_size, learning_rate=learning_rate, dropout_rate=dropout_rate ) # 提前停止回调 early_stop_callback = EarlyStopping( monitor='val_loss', patience=5, verbose=False, mode='min' ) # Optuna 剪枝回调 pruning_callback = PyTorchLightningPruningCallback(trial, monitor='val_loss') # 日志记录器 logger = TensorBoardLogger(save_dir=os.getcwd(), name=f"optuna_logs/trial_{trial.number}") # 创建训练器 trainer = pl.Trainer( max_epochs=10, callbacks=[early_stop_callback, pruning_callback], logger=logger, enable_progress_bar=False, enable_model_summary=False ) # 准备数据 train_loader, val_loader, test_loader = prepare_data() # 训练模型 trainer.fit(model, train_loader, val_loader) # 最终验证损失 return trainer.callback_metrics['val_loss'].item() |
另一个用于管理优化过程的 Optuna 关键函数是 run_optimization，在该函数中，我们根据前面优化函数中定义的规范，指示要运行的随机试验次数。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 |
def run_optimization(n_trials=20): pruner = optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=10) study = optuna.create_study(direction='minimize', pruner=pruner) study.optimize(objective, n_trials=n_trials) print("最佳试验：") trial = study.best_trial print(f" 值：{trial.value}") print(" 参数：") for key, value in trial.params.items(): print(f" {key}: {value}") return study |
在完成超参数优化过程并确定最佳模型配置后，需要另一个函数来接收这些结果，并在测试集上评估该模型的性能以进行最终验证。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
def test_best_model(study): # 获取最佳超参数 best_params = study.best_trial.params # 使用最佳超参数创建模型 model = MNISTClassifier( layer_1_size=best_params['layer_1_size'], layer_2_size=best_params['layer_2_size'], learning_rate=best_params['learning_rate'], dropout_rate=best_params['dropout_rate'] ) # 创建 Trainer 实例 trainer = pl.Trainer(max_epochs=10) # 准备数据 train_loader, val_loader, test_loader = prepare_data() # 使用最佳超参数训练模型 trainer.fit(model, train_loader, val_loader) # 使用测试数据测试模型 results = trainer.test(model, test_loader) return results |
现在我们有了所有需要的类和函数，最后通过一个将它们整合在一起的演示来完成。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
study = run_optimization(n_trials=5) # 可视化结果 try: # 绘制优化历史 optuna.visualization.plot_optimization_history(study) # 绘制参数重要性 optuna.visualization.plot_param_importances(study) # 绘制平行坐标图 optuna.visualization.plot_parallel_coordinate(study) except ImportError: print("可视化需要 plotly。请使用以下命令安装：pip install plotly") # 测试最佳模型 results = test_best_model(study) print(f"使用最佳超参数的测试结果：{results}") |
以下是工作流程的概要总结
- 运行一个 study（一组 Optuna 实验），指定试验次数
- 对训练过程中的优化过程进行一系列可视化
- 找到最佳模型后，将其暴露给测试数据以进一步评估
总结
本文介绍了如何将 PyTorch Lightning 和 Optuna 结合使用，为神经网络模型执行高效且有效的超参数优化。Optuna 提供了增强的模型调优算法，而 PyTorch Lightning 则构建在 PyTorch 之上，以更高级的方式进一步简化神经网络建模过程。
暂无评论。