PyTorch Lightning CLI 进阶指南：多模型与数据集的灵活配置-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00887/article/details/148360213

PyTorch Lightning CLI 进阶指南：多模型与数据集的灵活配置

pytorch-lightning Lightning-AI/pytorch-lightning: PyTorch Lightning 是一个轻量级的高级接口，用于简化 PyTorch 中深度学习模型的训练流程。它抽象出了繁杂的工程细节，使研究者能够专注于模型本身的逻辑和实验设计，同时仍能充分利用PyTorch底层的灵活性。项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning

前言

在深度学习项目开发过程中，随着项目复杂度增加，我们经常需要处理多个模型和数据集组合的情况。PyTorch Lightning 提供的 LightningCLI 工具能够极大简化这一过程，让开发者可以轻松通过命令行接口(CLI)实现模型与数据集的自由组合。本文将深入探讨如何利用这一功能提升开发效率。

为什么需要多模型与数据集组合

在项目初期，我们通常只处理一个模型和一个数据集。但随着项目发展，这种简单组合往往无法满足需求。例如：

需要比较不同模型在同一数据集上的表现
需要测试同一模型在不同数据集上的泛化能力
需要快速切换不同的优化器和学习率调度器

传统实现方式需要编写大量条件判断代码，不仅冗长而且难以维护：

# 传统方式需要大量条件判断
if args.model == "gan":
    model = GAN(args.feat_dim)
elif args.model == "transformer":
    model = Transformer(args.feat_dim)
...

PyTorch Lightning 的 LightningCLI 提供了一种更优雅的解决方案。

配置多个 LightningModule

基础配置方法

要支持多个模型，在实例化 LightningCLI 时不指定 model_class 参数：

from lightning.pytorch.cli import LightningCLI

class Model1(DemoModel):
    def configure_optimizers(self):
        print("使用 Model1")
        return super().configure_optimizers()

class Model2(DemoModel):
    def configure_optimizers(self):
        print("使用 Model2")
        return super().configure_optimizers()

cli = LightningCLI(datamodule_class=BoringDataModule)

使用技巧

可以通过命令行自由选择模型：

# 使用 Model1
python main.py fit --model Model1

# 使用 Model2
python main.py fit --model Model2

专业建议：可以使用 subclass_mode_model=True 参数限制只接受特定基类的子类作为模型，提高类型安全性。

配置多个 LightningDataModule

基础配置方法

类似地，要支持多个数据集模块，不指定 datamodule_class 参数：

class FakeDataset1(BoringDataModule):
    def train_dataloader(self):
        print("使用 FakeDataset1")
        return DataLoader(self.random_train)

class FakeDataset2(BoringDataModule):
    def train_dataloader(self):
        print("使用 FakeDataset2")
        return DataLoader(self.random_train)

cli = LightningCLI(DemoModel)

使用技巧

通过命令行选择数据集：

# 使用 FakeDataset1
python main.py fit --data FakeDataset1

# 使用 FakeDataset2
python main.py fit --data FakeDataset2

专业建议：同样可以使用 subclass_mode_data=True 限制数据集模块的类型。

优化器与调度器的灵活配置

标准优化器使用

PyTorch 内置优化器可以直接使用：

python main.py fit --optimizer AdamW

自定义优化器参数无需修改代码：

python main.py fit --optimizer SGD --optimizer.lr=0.01

自定义优化器

可以轻松集成自定义优化器：

class LitAdam(torch.optim.Adam):
    def step(self, closure):
        print("使用 LitAdam")
        super().step(closure)

学习率调度器

标准调度器同样开箱即用：

python main.py fit --optimizer=Adam --lr_scheduler CosineAnnealingLR

注意：必须同时指定 --optimizer 参数。

自定义调度器参数：

python main.py fit --optimizer=Adam --lr_scheduler=ReduceLROnPlateau --lr_scheduler.monitor=epoch

跨模块类引用

导入外部模块类

要使用其他模块中定义的类，只需导入相应模块：

import my_code.models  # 导入模型模块
import my_code.data_modules  # 导入数据模块
import my_code.optimizers  # 导入优化器模块

cli = LightningCLI()

完整路径引用

也可以通过完整路径引用未导入的类：

python main.py fit --model my_code.models.Model1

获取特定类帮助信息

当支持多个类选项时，可以通过以下方式获取特定类的详细参数帮助：

python main.py fit --model.help Model1
python main.py fit --data.help FakeDataset2
python main.py fit --optimizer.help Adagrad
python main.py fit --lr_scheduler.help StepLR