pytorch-lightning

1pip install lightning['extra']

bash
1conda install lightning -c conda-forge

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

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

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

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

python
1# main.py
2# ! pip install torchvision
3import torch, torch.nn as nn, torch.utils.data as data, torchvision as tv, torch.nn.functional as F
4import lightning as L
5
6# --------------------------------
7# 步骤1: 定义一个 LightningModule
8# --------------------------------
9# 一个 LightningModule（nn.Module 的子类）定义了一个完整的*系统*
10# （例如：一个 LLM、扩散模型、自编码器，或一个简单的图像分类器）。
11
12
13class LitAutoEncoder(L.LightningModule):
14    def __init__(
15
16self):
17        super().__init__()
18        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
19        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))
20
21    def forward(self, x):
22        # 在 lightning 中，forward 定义了预测/推理行为
23        embedding = self.encoder(x)
24        return embedding
25
26    def training_step(self, batch, batch_idx):
27        # training_step 定义了训练循环。它独立于 forward
28        x, _ = batch
29        x = x.view(x.size(0), -1)
30        z = self.encoder(x)
31        x_hat = self.decoder(z)
32        loss = F.mse_loss(x_hat, x)
33        self.log("train_loss", loss)
34        return loss
35
36    def configure_optimizers(self):
37        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
38        return optimizer
39
40
41# -------------------
42# 步骤2: 定义数据
43# -------------------
44dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
45train, val = data.random_split(dataset, [55000, 5000])
46
47# -------------------
48# 步骤3: 训练
49# -------------------
50autoencoder = LitAutoEncoder()
51trainer = L.Trainer()
52trainer.fit(autoencoder, data.DataLoader(train), data.DataLoader(val))

bash
1pip install torchvision
2python main.py

更改内容结果 Fabric 代码（复制我！）

更改内容	结果 Fabric 代码（复制我！）
_{diff 1+ import lightning as L 2 import torch; import torchvision as tv 3 4 dataset = tv.datasets.CIFAR10("data", download=True, 5 train=True, 6 transform=tv.transforms.ToTensor()) 7 8+ fabric = L.Fabric() 9+ fabric.launch() 10 11 model = tv.models.resnet18() 12 optimizer = torch.optim.SGD(model.parameters(), lr=0.001) 13- device = "cuda" if torch.cuda.is_available() else "cpu" 14- model.to(device) 15+ model, optimizer = fabric.setup(model, optimizer) 16 17 dataloader = torch.utils.data.DataLoader(dataset, batch_size=8) 18+ dataloader = fabric.setup_dataloaders(dataloader) 19 20 model.train() 21 num_epochs = 10 22 for epoch in range(num_epochs): 23 for batch in dataloader: 24 inputs, labels = batch 25- inputs, labels = inputs.to(device), labels.to(device) 26 optimizer.zero_grad() 27 outputs = model(inputs) 28 loss = torch.nn.functional.cross_entropy(outputs, labels) 29- loss.backward() 30+ fabric.backward(loss) 31 optimizer.step() 32 print(loss.data)}	_{Python 1import lightning as L 2import torch; import torchvision as tv 3 4dataset = tv.datasets.CIFAR10("data", download=True, 5 train=True, 6 transform=tv.transforms.ToTensor()) 7 8fabric = L.Fabric() 9fabric.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)}

_{diff
1+ import lightning as L
2 import torch; import torchvision as tv
3
4 dataset = tv.datasets.CIFAR10("data", download=True,
5 train=True,
6 transform=tv.transforms.ToTensor())
7
8+ fabric = L.Fabric()
9+ fabric.launch()
10
11 model = tv.models.resnet18()
12 optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
13- device = "cuda" if torch.cuda.is_available() else "cpu"
14- model.to(device)
15+ model, optimizer = fabric.setup(model, optimizer)
16
17 dataloader = torch.utils.data.DataLoader(dataset, batch_size=8)
18+ dataloader = fabric.setup_dataloaders(dataloader)
19
20 model.train()
21 num_epochs = 10
22 for epoch in range(num_epochs):
23 for batch in dataloader:
24 inputs, labels = batch
25- inputs, labels = inputs.to(device), labels.to(device)
26 optimizer.zero_grad()
27 outputs = model(inputs)
28 loss = torch.nn.functional.cross_entropy(outputs, labels)
29- loss.backward()
30+ fabric.backward(loss)
31 optimizer.step()
32 print(loss.data)}

_{Python
1import lightning as L
2import torch; import torchvision as tv
3
4dataset = tv.datasets.CIFAR10("data", download=True,
5 train=True,
6 transform=tv.transforms.ToTensor())
7
8fabric = L.Fabric()
9fabric.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)}

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


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 版本)

AI 奇想空间

AI 奇想空间

项目详情

安装 Lightning

安装带可选依赖项的版本

相关项目

exo

Conda 安装

安装稳定版本

安装最新版本

Lightning 有两个核心包

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

示例

Hello 简单模型

高级功能

相比未结构化的 PyTorch 的优势

Lightning Fabric: 专家控制

主要特性

示例

自监督学习

卷积架构

强化学习

GANs

经典机器学习

持续集成

*Codecov 超过 90%+，但构建延迟可能显示更少

社区

请求帮助

chatgpt-api

ChatGPT