PyTorch-Lightning 模板

中文版

简介

PyTorch-Lightning 是一个非常便捷的库，可以看作是对 PyTorch 的抽象和封装。它的优点在于可重用性强、易于维护、逻辑清晰等；缺点则是较为臃肿，需要花费一定时间去学习和理解。此外，由于它直接将模型与训练代码绑定在一起，因此并不适合包含多个模型和数据集文件的真实项目。同样的问题也存在于数据模块的设计中：DataLoader 和自定义 Dataset 之间的强耦合，导致代码在不同地方被不优雅地复制粘贴。

经过多次探索和实践，我总结出了以下模板，这些模板也可以被视为对 PyTorch-Lightning 的进一步抽象。在最初的版本中，所有模板内容都放在根目录下。然而，在使用了一个多月后，我发现针对不同类型项目的更细化模板能够显著提升编码效率。例如，分类任务和超分辨率任务都有一些固定的需求点，通过直接修改专门的模板，可以更快地实现项目代码，并减少一些本可以避免的 bug。

目前，由于这仍然是一个较新的库，所以只有这两个模板。不过，随着我在其他项目中的应用，未来还会添加更多专门化的模板。如果你已经使用这个模板完成了自己的任务（如 NLP、GAN、语音识别等），欢迎提交 PR，将你的模板整合到库中，以便更多人使用。如果你的任务尚未列入清单，可以从 classification 模板开始入手。由于大多数模板的底层逻辑和代码是相同的，因此可以非常快速地完成。

欢迎大家尝试这套代码风格。一旦习惯之后，复用起来会非常方便，也不容易陷入重复劳动的陷阱。关于 PyTorch-Lightning 的更详细解释和完整指南，请参阅知乎博客中的这篇文章。

文件结构

root-
	|-data
		|-__init__.py
		|-data_interface.py
		|-xxxdataset1.py
		|-xxxdataset2.py
		|-...
	|-model
		|-__init__.py
		|-model_interface.py
		|-xxxmodel1.py
		|-xxxmodel2.py
		|-...
	|-main.py
	|-utils.py

安装

无需安装。直接运行 git clone https://github.com/miracleyoo/pytorch-lightning-template.git 将其克隆到本地。根据你的问题类型选择如 classification 等模板，并将其复制到你的项目目录中。

结构说明

• 根目录下仅有 main.py 和 utils.py。前者是代码的入口，后者则是一些辅助文件。
• 在 data 和 model 文件夹中都包含一个 __init__.py 文件，用于将它们转换为包，从而简化导入操作。
• 在 data_interface 中创建一个 class DInterface(pl.LightningDataModule): 类，作为所有自定义 Dataset 文件的接口。相应的 Dataset 类会在 __init__() 函数中被导入。实例化工作在 setup() 方法中完成，并定义了 train_dataloader、val_dataloader 和 test_dataloader 方法。
• 同样地，在 model_interface 中创建一个 class MInterface(pl.LightningModule): 类，作为所有模型文件的接口。对应的模型类也会在 __init__() 函数中被导入。你只需要在接口中修改控制训练过程的函数，如 configure_optimizers、training_step 和 validation_step。一个接口适用于所有模型，差异部分则通过参数来处理。
• main.py 只负责以下任务：
  • 定义解析器，添加解析项。（注意：如果有某些参数需要在外部控制，比如通过命令行输入，可以直接在 main.py 文件中添加解析项。例如，有一个名为 random_arg 的字符串参数，可以在 main.py 中添加 parser.add_argument('--random_arg', default='test', type=str)。两个 Interface 类会自动选择并将这些参数传递给相应的模型或数据类。）
  • 选择所需的回调函数，如自动保存、早停和学习率调度器。
  • 实例化 MInterface、DInterface 和 Trainer。

完毕。

注意事项

需要注意的一点是，为了让 MInterface 和 DInterface 能够通过简单指定 --model_name 和 --dataset 参数，自动解析你新添加的模型和数据集，我们对模型/数据集文件采用蛇形命名法（如 standard_net.py），而类名则采用驼峰命名法（如 StandardNet）。

data 文件夹也同样遵循这一规则。

虽然这看似限制了你对模型和数据集的命名，但这样做也能使代码更易读、易理解。请务必注意这一点，以避免解析时出现问题。

引用

如果你使用了本模板并发现它对你的研究有所帮助，请考虑引用我们的论文：

@article{ZHANG2023126388,
title = {Neuromorphic high-frequency 3D dancing pose estimation in dynamic environment},
journal = {Neurocomputing},
volume = {547},
pages = {126388},
year = {2023},
issn = {0925-2312},
doi = {https://doi.org/10.1016/j.neucom.2023.126388},
url = {https://www.sciencedirect.com/science/article/pii/S0925231223005118},
author = {Zhongyang Zhang and Kaidong Chai and Haowen Yu and Ramzi Majaj and Francesca Walsh and Edward Wang and Upal Mahbub and Hava Siegelmann and Donghyun Kim and Tauhidur Rahman},
keywords = {Event Camera, Dynamic Vision Sensor, Neuromorphic Camera, Simulator, Dataset, Deep Learning, Human Pose Estimation, 3D Human Pose Estimation, Technology-Mediated Dancing},
}

@InProceedings{Zhang_2022_WACV,
    author    = {Zhang, Zhongyang and Xu, Zhiyang and Ahmed, Zia and Salekin, Asif and Rahman, Tauhidur},
    title     = {Hyperspectral Image Super-Resolution in Arbitrary Input-Output Band Settings},
    booktitle = {Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV) Workshops},
    month     = {January},
    year      = {2022},
    pages     = {749-759}
}

PyTorch-Lightning-Template 快速上手指南

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS 或 Windows
• Python 版本：建议 Python 3.8 及以上
• 前置依赖：
  • PyTorch
  • PyTorch Lightning (pytorch-lightning)
  • 其他常规深度学习库（如 numpy, tqdm 等）

提示：本项目本身无需通过 pip 安装，它是一个代码模板库。请确保已安装好 PyTorch 和 PyTorch Lightning 基础环境。国内用户可使用清华源加速安装基础依赖：

pip install torch torchvision torchaudio pytorch-lightning -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

本项目采用克隆源码并复制模板的方式使用，具体步骤如下：

1. 克隆仓库 将项目克隆到本地任意位置：

   git clone https://github.com/miracleyoo/pytorch-lightning-template.git
   

   国内加速：如果 GitHub 访问缓慢，可使用国内镜像（如 Gitee 镜像，若有）或通过代理加速克隆。

2. 选择并复制模板 进入克隆后的目录，根据您的任务类型（例如 classification 分类任务），将对应的模板文件夹复制到您的新项目目录中。

   假设您要创建一个分类项目：

   # 进入模板库根目录
   cd pytorch-lightning-template
   
   # 将 classification 模板复制到您的新项目文件夹 (例如 my_project)
   cp -r templates/classification ../my_project
   
   # 或者手动复制您需要的特定模板文件结构
   

   注：根据当前版本，主要提供 classification（分类）和 super-resolution（超分）等专用模板。若您的任务不在列表中，建议从 classification 模板开始修改。

基本使用

复制模板后，您的项目结构应包含 data/, model/, main.py 和 utils.py。以下是快速启动流程：

1. 理解核心结构

• main.py：程序入口。负责解析命令行参数、配置回调函数（如自动保存、早停、学习率调度）以及实例化 Trainer。
• data/：数据模块。
  • data_interface.py：定义 DInterface 类，统一处理 DataLoader 的构建。
  • xxxdataset.py：具体的 Dataset 实现文件。
• model/：模型模块。
  • model_interface.py：定义 MInterface 类，统一处理训练步 (training_step)、验证步及优化器配置。
  • xxxmodel.py：具体的网络模型实现文件。

2. 命名规范（重要）

为了让 main.py 自动识别并加载您新增的模型或数据集，必须严格遵守以下命名规则：

• 文件名：使用蛇形命名法（snake_case），例如 standard_net.py。
• 类名：使用驼峰命名法（CamelCase），且内容需与文件名对应，例如 StandardNet。
• 调用方式：运行时通过 --model_name standard_net 和 --dataset xxxdataset 参数自动匹配。

3. 运行示例

假设您已经在 model/ 下添加了 my_model.py (类名 MyModel)，在 data/ 下添加了 my_data.py (类名 MyData)。

在项目根目录执行：

python main.py --model_name my_model --dataset my_data --gpus 1 --max_epochs 50

自定义参数扩展： 如果需要传递额外参数（如随机种子或特定超参），可直接在 main.py 的 parser 部分添加：

# 在 main.py 中
parser.add_argument('--random_arg', default='test', type=str)

该参数会自动透传给对应的 Model 或 Data 接口类。

4. 开发新任务

• 新增模型：在 model/ 目录下新建 .py 文件，继承 pl.LightningModule 实现网络结构，并在 model_interface.py 中按需调整逻辑（通常只需关注 configure_optimizers 等核心训练流程）。
• 新增数据：在 data/ 目录下新建 .py 文件，继承 torch.utils.data.Dataset，并在 data_interface.py 的 setup() 中完成实例化。

通过以上步骤，您可以基于此模板快速搭建结构清晰、易于维护的 PyTorch Lightning 项目。