使用Ray Tune+PyTorch Lightning做HPO时，为何出现任务未完成及NameError？

错误原因 #

NameError: name 'LightningMNISTClassifier' is not defined的核心原因是代码中引用的自定义PyTorch Lightning模型类未被定义或导入。Ray官方文档的HPO示例默认你已经实现了这个针对MNIST任务的模型类，但你可能只复制了超参数优化的核心逻辑，遗漏了模型本身的代码实现。

解决步骤 #

1. 补全LightningMNISTClassifier类的实现 #

将以下适配MNIST任务的PyTorch Lightning模型代码添加到你的脚本中（放在调参逻辑之前）：

import pytorch_lightning as pl
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms

class LightningMNISTClassifier(pl.LightningModule):
    def __init__(self, hidden_size=64, learning_rate=1e-3):
        super().__init__()
        self.save_hyperparameters()
        
        self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_size)
        self.l2 = torch.nn.Linear(self.hparams.hidden_size, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = torch.relu(self.l1(x))
        x = self.l2(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)
        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 = (preds == y).float().mean()
        self.log("val_loss", loss)
        self.log("val_acc", acc)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

    def prepare_data(self):
        datasets.MNIST(root="./data", train=True, download=True)
        datasets.MNIST(root="./data", train=False, download=True)

    def setup(self, stage=None):
        if stage == "fit" or stage is None:
            mnist_full = datasets.MNIST(root="./data", train=True, transform=transforms.ToTensor())
            self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=32)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=32)

2. 确保依赖库安装完整 #

运行代码前执行以下命令安装所需依赖：

pip install ray[tune] pytorch-lightning torchvision

3. 整合Ray Tune超参数优化逻辑 #

将你复制的Ray HPO代码与上述模型类放在同一脚本中，确保模型类定义在调参函数之前。以下是典型的整合示例：

from ray import tune
from ray.tune.schedulers import ASHAScheduler
from ray.tune.integration.pytorch_lightning import TuneReportCallback

def train_mnist(config):
    model = LightningMNISTClassifier(**config)
    trainer = pl.Trainer(
        max_epochs=10,
        callbacks=[TuneReportCallback({"val_acc": "val_acc"})],
        logger=False,
    )
    trainer.fit(model)

search_space = {
    "hidden_size": tune.choice([32, 64, 128]),
    "learning_rate": tune.loguniform(1e-4, 1e-2),
}

scheduler = ASHAScheduler(
    max_t=10,
    grace_period=1,
    reduction_factor=2,
)

analysis = tune.run(
    train_mnist,
    metric="val_acc",
    mode="max",
    config=search_space,
    scheduler=scheduler,
    num_samples=10,
    resources_per_trial={"cpu": 2, "gpu": 0},
)

print("Best hyperparameters found: ", analysis.best_config)

关键注意事项 #

• 若模型类定义在其他模块中，需通过from your_module import LightningMNISTClassifier正确导入。
• 若你的任务不是MNIST，只需修改模型类的网络结构、数据处理逻辑，保持与PyTorch Lightning的接口规范一致即可。
• TuneReportCallback负责将PyTorch Lightning的训练指标上报给Ray Tune，用于超参数性能评估。

内容的提问来源于stack exchange，提问作者Slowat_Kela