Lightning-Hydra-Template: 一个用户友好的机器学习实验模板

Lightning-Hydra-Template简介

Lightning-Hydra-Template是一个基于PyTorch Lightning和Hydra的机器学习项目模板,旨在为研究人员和工程师提供一个灵活、可扩展且易于使用的实验框架。该模板集成了多种常用的机器学习工具和最佳实践,可以大大简化实验流程,提高代码的复用性和可维护性。

主要特性

快速实验: 利用Hydra的命令行超能力,可以轻松地进行参数覆盖和多次运行。
最小样板代码: 通过配置实例化自动化管道,减少重复代码。
灵活配置: 主配置文件和实验配置文件允许您指定默认训练配置并轻松覆盖超参数。
工作流简化: 整个工作流程简化为4个简单步骤。
实验跟踪: 支持多种流行的日志框架,如Tensorboard、Weights & Biases等。
日志管理: 所有日志(检查点、配置等)都存储在动态生成的文件夹结构中。
超参数搜索: 使用Hydra插件(如Optuna Sweeper)可以轻松进行超参数搜索。
测试与持续集成: 包含通用的烟雾测试,并通过Github Actions实现自动测试和代码检查。

为什么选择Lightning-Hydra-Template?

节省样板代码: 轻松添加新模型、数据集、任务和实验,并在不同的加速器(如多GPU、TPU或SLURM集群)上进行训练。
教育价值: 代码有详细注释,可以作为学习资源使用。
可重用性: 收集了有用的MLOps工具、配置和代码片段,可以作为各种实用程序的参考。

然而,也需要注意该模板的一些局限性:

Lightning和Hydra仍在不断发展,可能会出现一些问题。
模板主要针对简单的Lightning训练场景,可能需要调整以适应不同的用例。
不支持恢复基于Hydra的多次运行或超参数搜索。

核心技术

Lightning-Hydra-Template主要基于两项核心技术:

PyTorch Lightning: 一个轻量级的PyTorch包装器,用于高性能AI研究。它提供了一个组织PyTorch代码的框架,简化了训练循环和分布式训练等复杂任务。
Hydra: 一个用于优雅配置复杂应用程序的框架。其关键特性是能够通过组合动态创建分层配置,并通过配置文件和命令行进行覆盖。

这两项技术的结合为机器学习实验提供了强大而灵活的基础。

项目结构

Lightning-Hydra-Template提供了一个清晰的项目结构,有助于组织和管理机器学习项目:

├── configs                   <- Hydra配置文件
├── data                      <- 项目数据
├── logs                      <- 日志文件
├── notebooks                 <- Jupyter笔记本
├── scripts                   <- Shell脚本
├── src                       <- 源代码
│   ├── data                  <- 数据处理脚本
│   ├── models                <- 模型定义脚本
│   ├── utils                 <- 实用工具脚本
│   ├── eval.py               <- 评估脚本
│   └── train.py              <- 训练脚本
├── tests                     <- 测试文件
└── ...                       <- 其他配置文件

这种结构使得项目的各个组件井然有序,便于开发和维护。

快速开始

要开始使用Lightning-Hydra-Template,您可以按照以下步骤操作:

克隆项目仓库:

git clone https://github.com/ashleve/lightning-hydra-template
cd lightning-hydra-template

创建并激活conda环境(可选):

conda create -n myenv python=3.9
conda activate myenv

安装依赖:
```
pip install -r requirements.txt
```
运行示例训练:
```
python src/train.py
```

模板包含一个MNIST分类的示例。运行上述命令后,您应该能看到训练过程的输出。

训练过程示例

超能力:灵活的配置和运行

Lightning-Hydra-Template提供了多种"超能力",使得实验配置和运行变得异常灵活:

从命令行覆盖任何配置参数:

python train.py trainer.max_epochs=20 model.optimizer.lr=1e-4

在不同硬件上训练:

# CPU训练
python train.py trainer=cpu

# 单GPU训练
python train.py trainer=gpu

# 多GPU分布式数据并行(DDP)训练
python train.py trainer=ddp trainer.devices=4

使用混合精度训练:

python train.py trainer=gpu +trainer.precision=16

使用不同的日志工具:
```
python train.py logger=wandb
```
使用预定义的实验配置:
```
python train.py experiment=example
```
添加回调函数:
```
python train.py callbacks=default
```
轻松调试:
```
python train.py debug=default
```

从检查点恢复训练:

python train.py ckpt_path="/path/to/ckpt/name.ckpt"

创建超参数搜索:

python train.py -m data.batch_size=32,64,128 model.lr=0.001,0.0005

使用Optuna进行超参数优化:

python train.py -m hparams_search=mnist_optuna experiment=example

这些功能大大提高了实验的灵活性和效率,使研究人员能够快速尝试不同的配置和方法。

工作流程

使用Lightning-Hydra-Template的基本工作流程如下:

编写PyTorch Lightning模块(参考src/models/mnist_module.py)
编写PyTorch Lightning数据模块(参考src/data/mnist_datamodule.py)
编写实验配置,包含模型和数据模块的路径

使用选定的实验配置运行训练:

python src/train.py experiment=experiment_name.yaml

对于更复杂的实验设计,例如研究批量大小对准确率的影响,可以按以下方式进行:

执行多次运行,使用标签来标识:

python train.py -m logger=csv data.batch_size=16,32,64,128 tags=["batch_size_exp"]

编写脚本或笔记本,搜索logs/文件夹并检索包含给定标签的CSV日志,然后绘制结果。

日志和实验跟踪

Lightning-Hydra-Template支持多种流行的日志框架,如Weights & Biases、Neptune、Comet、MLFlow和Tensorboard。这些工具可以帮助您跟踪超参数和输出指标,并允许比较和可视化结果。

要使用特定的日志工具,只需完成其在configs/logger中的配置,然后运行:

python train.py logger=logger_name

此外,Lightning提供了一种从LightningModule内部记录自定义指标的便捷方法:

self.log("train/loss", loss)

这种方式可以帮助组织和分类不同的指标,使得在UI中更易于查看和分析。

测试和持续集成

模板附带了使用pytest实现的通用测试。这些测试主要用于验证某些命令的执行不会导致异常,可以帮助加速开发过程。

运行测试的方法:

# 运行所有测试
pytest

# 运行特定文件的测试
pytest tests/test_train.py

# 运行除了标记为慢速的所有测试
pytest -k "not slow"

此外,模板还包含了在Github Actions中实现的CI工作流:

使用pytest运行所有测试
在主分支上对所有文件运行预提交检查
在拉取请求上仅对修改的文件运行预提交检查

这些CI工作流可以帮助保持代码质量,并在引入新功能或修复时捕获潜在问题。

分布式训练

Lightning支持多种分布式训练方法。最常见的是DDP(Distributed Data Parallel),它为每个GPU生成单独的进程并在它们之间平均梯度。

要在4个GPU上运行DDP的MNIST示例,可以使用以下命令:

python train.py trainer=ddp

需要注意的是,使用DDP时需要小心编写模型,以确保正确的行为。

最佳实践

Lightning-Hydra-Template还提供了一些最佳实践建议,以提高项目的质量和可维护性:

使用Miniconda: 它比完整的Anaconda环境更轻量,足以满足大多数需求。
使用自动代码格式化: 使用pre-commit钩子来标准化项目的代码格式,节省精力。
在.env文件中设置私有环境变量: 系统特定变量不应该被版本控制。
使用'/'字符命名指标: 这有助于日志工具将指标组织到不同的部分。
使用torchmetrics: 使用官方的torchmetrics库确保正确计算指标,特别是在多GPU训练中。
遵循PyTorch Lightning风格指南: 这有助于保持代码的一致性和可读性。
使用DVC版本控制数据和模型: 对于大文件,如数据或训练好的ML模型,使用DVC进行版本控制。
支持将项目安装为包: 这允许其他人在自己的项目中轻松使用您的模块。
将本地配置保持在代码版本控制之外: 对于用户/机器/安装特定的配置,使用不被Git跟踪的本地配置文件。

通过遵循这些最佳实践,您可以创建更加健壮、可维护和可复用的机器学习项目。

结论

Lightning-Hydra-Template为机器学习实验提供了一个强大而灵活的框架。通过集成PyTorch Lightning和Hydra,它简化了实验配置、训练过程和结果跟踪,同时保持了高度的可定制性。无论您是研究人员还是工程师,这个模板都可以帮助您更高效地进行机器学习实验,专注于核心算法和模型开发,而不是繁琐的基础设施工作。

尽管该模板非常强大,但它仍在不断发展中。社区的贡献和反馈对于改进和扩展其功能至关重要。如果您在使用过程中遇到任何问题或有改进建议,欢迎在GitHub上提出issue或贡献代码。

通过使用Lightning-Hydra-Template,您可以将更多精力集中在创新和突破性研究上,而不是重复造轮子。让我们一起推动机器学习研究的边界,创造更多令人兴奋的成果!