PyTorch Lightning快速学习教程一:快速训练一个基础模型

粉丝量突破1200了!找到了喜欢的岗位,毕业上班刚好也有20天,为了督促自己终身学习的态度,继续开始坚持写写博客,沉淀并总结知识!

介绍:PyTorch Lightning是针对科研人员、机器学习开发者专门设计的,能够快速复用代码的一个工具,避免了因为每次都编写相似的代码而带来的时间成本。其可以理解为,lightning设计了一个,能够快速搭建训练验证测试模型的整套代码模板,我们只需要编写设计需要的模型、超参数、优化器等,直接套进去即可。lightning的优势在于:灵活性高、可读性强、支持多卡训练、内置测试、内置日志等。

前置掌握知识:Python和PyTorch的使用

链接:https://lightning.ai/

快速安装:pip install lightning

1.添加依赖包

需要添加相应的依赖,包括os,torch工具包,torch数据载入等依赖

python 复制代码
import os		
import torch
from torch import nn
import torch.nn.functional as F
from torchvision import transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
import lightning.pytorch as pl

2.定义模型

PyTorch定义模型案例如下,定义好了方便后续的调用

python 复制代码
class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
    def forward(self, x):
        return self.l1(x)	# 全连接 激活 全连接

class Decoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))
    def forward(self, x):
        return self.l1(x)	# 全连接 激活 全连接

3.定义网络架构

定义网络模型,自定义模型名字,并继承lightning.pytorch.LightningModule类,如下代码

  • training_step定义了与nn.Module之间交互

  • configure_optimizers为模型定义优化器

python 复制代码
class LitAutoEncoder(pl.LightningModule):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def training_step(self, batch, batch_idx):
        # training_step defines the train loop.
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        return loss

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

4.定义训练集

定义DataLoader,这一点跟PyTorch调模型的流程一样,如下调用了MNIST公开数据集

python 复制代码
dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
train_loader = DataLoader(dataset)

5.训练数据

使用Lightning来处理所有的训练,如下代码。

python 复制代码
# model 模型
autoencoder = LitAutoEncoder(Encoder(), Decoder())

# train model 训练
trainer = pl.Trainer()
trainer.fit(model=autoencoder, train_dataloaders=train_loader)

一般的训练过程,需要设计如下代码,进行遍历和循环训练,Lightning会消除这些繁琐的过程,使用Lightning,可以将所有这些技术混合在一起,而无需每次都重写一个新的循环。

python 复制代码
autoencoder = LitAutoEncoder(Encoder(), Decoder())
optimizer = autoencoder.configure_optimizers()

for batch_idx, batch in enumerate(train_loader):
    loss = autoencoder.training_step(batch, batch_idx)

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

完整代码

python 复制代码
# coding:utf-8
import torch, torch.nn as nn, torch.utils.data as data, torchvision as tv, torch.nn.functional as F
import lightning as L

# --------------------------------
# Step 1: 定义一个 LightningModule
# --------------------------------
# A LightningModule (nn.Module subclass) defines a full *system*
# (例如: an LLM, diffusion model, autoencoder, or simple image classifier).


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):
        # forward 定义了一次 预测/推理 行为
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        # training_step 定义了一次训练的迭代, 和forward相互独立
        x, y = 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

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

# -------------------
# Step 3: 开始训练
# -------------------
autoencoder = LitAutoEncoder()
trainer = L.Trainer(accelerator="gpu")	
trainer.fit(autoencoder, data.DataLoader(train,batch_size=128), data.DataLoader(val))
相关推荐
行云流水剑10 分钟前
【学习记录】如何使用 Python 提取 PDF 文件中的内容
python·学习·pdf
算家计算13 分钟前
字节开源代码模型——Seed-Coder 本地部署教程,模型自驱动数据筛选,让每行代码都精准落位!
人工智能·开源
伪_装20 分钟前
大语言模型(LLM)面试问题集
人工智能·语言模型·自然语言处理
gs8014027 分钟前
Tavily 技术详解:为大模型提供实时搜索增强的利器
人工智能·rag
music&movie42 分钟前
算法工程师认知水平要求总结
人工智能·算法
虾球xz1 小时前
CppCon 2015 学习:CLANG/C2 for Windows
开发语言·c++·windows·学习
量子位1 小时前
苹果炮轰推理模型全是假思考!4 个游戏戳破神话,o3/DeepSeek 高难度全崩溃
人工智能·deepseek
黑鹿0221 小时前
机器学习基础(四) 决策树
人工智能·决策树·机器学习
Fxrain1 小时前
[深度学习]搭建开发平台及Tensor基础
人工智能·深度学习
szxinmai主板定制专家1 小时前
【飞腾AI加固服务器】全国产化飞腾+昇腾310+PCIe Switch的AI大模型服务器解决方案
运维·服务器·arm开发·人工智能·fpga开发