Loading...

pytorch-lightning - 开源项目详情

开源项目基础框架本页

基础框架

pytorch-lightning

Created by

Lightning-AI

Pretrain, finetune ANY AI model of ANY size on 1 or 10,000+ GPUs with zero code changes.

pytorch-lightning 预览图

项目详情

深入了解 pytorch-lightning 的功能与特性

安装 Lightning

从 PyPI 简单安装

bash

pip install lightning

其他安装选项

安装带可选依赖项的版本

bash

pip install lightning['extra']

Conda 安装

bash

conda install lightning -c conda-forge

安装稳定版本

从源码安装未来发布的稳定版本

bash

pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/release/stable.zip -U

安装最新版本

从源码安装 nightly 版本（不保证稳定性）

bash

pip install https://github.com/Lightning-AI/lightning/archive/refs/heads/master.zip -U

或从测试 PyPI 安装

bash

pip install -iU https://test.pypi.org/simple/ pytorch-lightning

Lightning 有两个核心包

PyTorch Lightning: 大规模训练和部署 PyTorch。
Lightning Fabric: 专家控制。

Lightning 让你可以精细控制在 PyTorch 上增加的抽象层次。

PyTorch Lightning: 大规模训练和部署 PyTorch

PyTorch Lightning 只是更有组织的 PyTorch - Lightning 解耦了 PyTorch 代码，将科学与工程分离。

PT to PL

31.1kStars

3.7kForks

257Watch

免责声明：本站大资源来自网络收集整理，小部分资源来自原创，如有侵权等，请联系处理。

相关项目

发现更多类似的优秀工具

exo

exo

Run frontier AI locally.

chatgpt-api

chatgpt-api

Node.js client for the official ChatGPT API. 🔥

#ai#chatbot#chatgpt

31.1kStars

3.7kForks

257Watch

任务	描述	运行
Hello world	预训练 - Hello world 示例
图像分割	微调 - 使用 ResNet-50 模型进行图像分割
文本分类	微调 - 文本分类器（BERT 模型）
文本摘要	微调 - 文本摘要（Hugging Face transformer 模型）
音频生成	微调 - 音频生成（transformer 模型）

python

# main.py
# ! pip install torchvision
import torch, torch.nn as nn, torch.utils.data as data, torchvision as tv, torch.nn.functional as F
import lightning as L

# --------------------------------
# 步骤1: 定义一个 LightningModule
# --------------------------------
# 一个 LightningModule（nn.Module 的子类）定义了一个完整的*系统*
# （例如：一个 LLM、扩散模型、自编码器，或一个简单的图像分类器）。


class LitAutoEncoder(L.LightningModule):
    def __init__(

self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))

    def forward(self, x):
        # 在 lightning 中，forward 定义了预测/推理行为
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        # training_step 定义了训练循环。它独立于 forward
        x, _ = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer


# -------------------
# 步骤2: 定义数据
# -------------------
dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
train, val = data.random_split(dataset, [55000, 5000])

# -------------------
# 步骤3: 训练
# -------------------
autoencoder = LitAutoEncoder()
trainer = L.Trainer()
trainer.fit(autoencoder, data.DataLoader(train), data.DataLoader(val))

在成千上万的GPU上训练而无需更改代码

python

# 8 个 GPU
# 无需代码更改
trainer = Trainer(accelerator="gpu", devices=8)

# 256 个 GPU
trainer = Trainer(accelerator="gpu", devices=8, num_nodes=32)

实验管理器

python

from lightning import loggers

# 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())

# neptune
trainer = Trainer(logger=loggers.NeptuneLogger())

# 还有很多

导出到 ONNX（生产用途）

python

# 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 代码（复制我！）

_{diff
+ 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)}

_{Python
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 甚至多节点训练

python

# 使用你的可用硬件
# 无需代码更改
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）和混合精度

python

# 使用最先进的分布式训练技术
fabric = Fabric(strategy="ddp")
fabric = Fabric(strategy="deepspeed")
fabric = Fabric(strategy="fsdp")

# 切换精度
fabric = Fabric(precision="16-mixed")
fabric = Fabric(precision="64")

使用 Fabric 原语构建你自己的自定义 Trainer 进行训练检查点记录、日志记录等

python


python
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()

你可以在我们的示例中找到更详细的示例

当前构建状态

系统 / PyTorch 版本	1.13	2.0	2.1
Linux py3.9 [GPUs]
Linux py3.9 [TPUs]
Linux (多个 Python 版本)
OSX (多个 Python 版本)
Windows (多个 Python 版本)