Lightning-AI/pytorch-lightning
GitHub: Lightning-AI/pytorch-lightning
PyTorch Lightning是一个深度学习框架，通过抽象训练工程代码让研究者专注于模型逻辑，支持从单卡到数千GPU的无缝扩展。
Stars: 31121 | Forks: 3719
**用于预训练和微调AI模型的深度学习框架。** **需要部署模型？** 使用 [LitServe](https://github.com/Lightning-AI/litserve?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) 用纯 Python 构建自定义推理服务器。
快速开始 • 示例 • PyTorch Lightning • Fabric • Lightning Cloud • 社区 • 文档
[](https://pypi.org/project/pytorch-lightning/) [](https://badge.fury.io/py/pytorch-lightning) [](https://pepy.tech/project/pytorch-lightning) [](https://anaconda.org/conda-forge/lightning) [](https://codecov.io/gh/Lightning-AI/pytorch-lightning) [](https://discord.gg/VptPCZkGNa)  [](https://github.com/Lightning-AI/pytorch-lightning/blob/master/LICENSE)高级安装选项
#### 安装可选依赖 ``` pip install lightning['extra'] ``` #### Conda ``` conda install lightning -c conda-forge ``` #### 安装稳定版本 从源码安装未来版本 ``` pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/release/stable.zip -U ``` #### 安装最新版本 从源码安装每日构建版本（无保证） ``` pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/master.zip -U ``` 或从 testing PyPI 安装 ``` pip install -iU https://test.pypi.org/simple/ pytorch-lightning ```
|
| [图像分类](https://lightning.ai/lightning-ai/studios/image-classification-with-pytorch-lightning?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) | 微调 - ResNet-34 模型对汽车图像进行分类 | |
| [图像分割](https://lightning.ai/lightning-ai/studios/image-segmentation-with-pytorch-lightning?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) | 微调 - ResNet-50 模型分割图像 | |
| [目标检测](https://lightning.ai/lightning-ai/studios/object-detection-with-pytorch-lightning?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) | 微调 - Faster R-CNN 模型检测目标 | |
| [文本分类](https://lightning.ai/lightning-ai/studios/text-classification-with-pytorch-lightning?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) | 微调 - 文本分类器（BERT 模型） | |
| [文本摘要](https://lightning.ai/lightning-ai/studios/text-summarization-with-pytorch-lightning?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) | 微调 - 文本摘要（Hugging Face transformer 模型） | |
| [音频生成](https://lightning.ai/lightning-ai/studios/finetune-a-personal-ai-music-generator?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) | 微调 - 音频生成器（transformer 模型） | |
| [LLM 微调](https://lightning.ai/lightning-ai/studios/finetune-an-llm-with-pytorch-lightning?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) | 微调 - LLM（Meta Llama 3.1 8B） | |
| [图像生成](https://lightning.ai/lightning-ai/studios/train-a-diffusion-model-with-pytorch-lightning?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) | 预训练 - 图像生成器（diffusion 模型） | |
| [推荐系统](https://lightning.ai/lightning-ai/studios/recommendation-system-with-pytorch-lightning?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) | 训练 - 推荐系统（因式分解和嵌入） | |
| [时间序列预测](https://lightning.ai/lightning-ai/studios/time-series-forecasting-with-pytorch-lightning?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme) | 训练 - 使用 LSTM 进行时间序列预测 | |
## 高级功能
Lightning 拥有超过 [40+ 高级功能](https://lightning.ai/docs/pytorch/stable/common/trainer.html?utm_source=ptl_readme&utm_medium=referral&utm_campaign=ptl_readme#trainer-flags)，专为大规模专业 AI 研究设计。
以下是一些示例：
无需更改代码即可在数千个 GPU 上训练
``` # 8 个 GPU # 无需代码更改 trainer = Trainer(accelerator="gpu", devices=8) # 256 个 GPU trainer = Trainer(accelerator="gpu", devices=8, num_nodes=32) ```无需更改代码即可在其他加速器（如 TPU）上训练
``` # 无需代码更改 trainer = Trainer(accelerator="tpu", devices=8) ```16 位精度
``` # 无需代码更改 trainer = Trainer(precision=16) ```实验管理器
``` from lightning import loggers # litlogger trainer = Trainer(logger=LitLogger()) # tensorboard trainer = Trainer(logger=TensorBoardLogger("logs/")) # weights and biases trainer = Trainer(logger=loggers.WandbLogger()) # comet trainer = Trainer(logger=loggers.CometLogger()) # mlflow trainer = Trainer(logger=loggers.MLFlowLogger()) # ... 以及数十个更多 ```早停
``` es = EarlyStopping(monitor="val_loss") trainer = Trainer(callbacks=[es]) ```检查点保存
``` checkpointing = ModelCheckpoint(monitor="val_loss") trainer = Trainer(callbacks=[checkpointing]) ```导出到 torchscript (JIT)（生产使用）
``` # torchscript autoencoder = LitAutoEncoder() torch.jit.save(autoencoder.to_torchscript(), "model.pt") ```导出到 ONNX（生产使用）
``` # onnx with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as tmpfile: autoencoder = LitAutoEncoder() input_sample = torch.randn((1, 64)) autoencoder.to_onnx(tmpfile.name, input_sample, export_params=True) os.path.isfile(tmpfile.name) ```| 需要修改的内容 | 结果 Fabric 代码（复制我的！） |
|---|---|
| ``` + import lightning as L import torch; import torchvision as tv dataset = tv.datasets.CIFAR10("data", download=True, train=True, transform=tv.transforms.ToTensor()) + fabric = L.Fabric() + fabric.launch() model = tv.models.resnet18() optimizer = torch.optim.SGD(model.parameters(), lr=0.001) - device = "cuda" if torch.cuda.is_available() else "cpu" - model.to(device) + model, optimizer = fabric.setup(model, optimizer) dataloader = torch.utils.data.DataLoader(dataset, batch_size=8) + dataloader = fabric.setup_dataloaders(dataloader) model.train() num_epochs = 10 for epoch in range(num_epochs): for batch in dataloader: inputs, labels = batch - inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = torch.nn.functional.cross_entropy(outputs, labels) - loss.backward() + fabric.backward(loss) optimizer.step() print(loss.data) ``` | ``` import lightning as L import torch; import torchvision as tv dataset = tv.datasets.CIFAR10("data", download=True, train=True, transform=tv.transforms.ToTensor()) fabric = L.Fabric() fabric.launch() model = tv.models.resnet18() optimizer = torch.optim.SGD(model.parameters(), lr=0.001) model, optimizer = fabric.setup(model, optimizer) dataloader = torch.utils.data.DataLoader(dataset, batch_size=8) dataloader = fabric.setup_dataloaders(dataloader) model.train() num_epochs = 10 for epoch in range(num_epochs): for batch in dataloader: inputs, labels = batch optimizer.zero_grad() outputs = model(inputs) loss = torch.nn.functional.cross_entropy(outputs, labels) fabric.backward(loss) optimizer.step() print(loss.data) ``` |
轻松切换从 CPU 到 GPU（Apple Silicon、CUDA…）、TPU、多 GPU 甚至多节点训练
``` # 使用您可用的硬件 # 无需代码更改 fabric = Fabric() # 在 GPU 上运行（CUDA 或 MPS） fabric = Fabric(accelerator="gpu") # 8 个 GPU fabric = Fabric(accelerator="gpu", devices=8) # 256 个 GPU，多节点 fabric = Fabric(accelerator="gpu", devices=8, num_nodes=32) # 在 TPU 上运行 fabric = Fabric(accelerator="tpu") ```开箱即用使用最先进的分布式训练策略（DDP、FSDP、DeepSpeed）和混合精度
``` # 使用最先进的分布式训练技术 fabric = Fabric(strategy="ddp") fabric = Fabric(strategy="deepspeed") fabric = Fabric(strategy="fsdp") # 切换精度 fabric = Fabric(precision="16-mixed") fabric = Fabric(precision="64") ```所有设备逻辑样板代码都为您处理好了
``` # no more of this! - model.to(device) - batch.to(device) ```使用 Fabric 原语构建您自己的自定义 Trainer，用于训练检查点、日志记录等
``` import lightning as L class MyCustomTrainer: def __init__(self, accelerator="auto", strategy="auto", devices="auto", precision="32-true"): self.fabric = L.Fabric(accelerator=accelerator, strategy=strategy, devices=devices, precision=precision) def fit(self, model, optimizer, dataloader, max_epochs): self.fabric.launch() model, optimizer = self.fabric.setup(model, optimizer) dataloader = self.fabric.setup_dataloaders(dataloader) model.train() for epoch in range(max_epochs): for batch in dataloader: input, target = batch optimizer.zero_grad() output = model(input) loss = loss_fn(output, target) self.fabric.backward(loss) optimizer.step() ``` 您可以在我们的[示例](当前构建状态
标签：AI模型部署, Apex, CNCF毕业项目, GPU计算, Lightning, Python, PyTorch, Vectored Exception Handling, 人工智能框架, 凭据扫描, 分布式训练, 多GPU训练, 大规模训练, 并行计算, 微调, 无后门, 机器学习, 模型优化, 模型训练, 深度学习, 深度学习框架, 神经网络, 训练框架, 逆向工具, 预训练, 高性能计算