粉絲量突破1200了！找到了喜歡的崗位，畢業(yè)上班剛好也有20天，為了督促自己終身學(xué)習(xí)的態(tài)度，繼續(xù)開始堅(jiān)持寫寫博客，沉淀并總結(jié)知識(shí)！
介紹：PyTorch Lightning是針對(duì)科研人員、機(jī)器學(xué)習(xí)開發(fā)者專門設(shè)計(jì)的，能夠快速?gòu)?fù)用代碼的一個(gè)工具，避免了因?yàn)槊看味季帉懴嗨频拇a而帶來(lái)的時(shí)間成本。其可以理解為，lightning設(shè)計(jì)了一個(gè)，能夠快速搭建訓(xùn)練驗(yàn)證測(cè)試模型的整套代碼模板，我們只需要編寫設(shè)計(jì)需要的模型、超參數(shù)、優(yōu)化器等，直接套進(jìn)去即可。lightning的優(yōu)勢(shì)在于：靈活性高、可讀性強(qiáng)、支持多卡訓(xùn)練、內(nèi)置測(cè)試、內(nèi)置日志等。

前置掌握知識(shí)：Python和PyTorch的使用

鏈接：https://lightning.ai/

快速安裝：pip install lightning

1.添加依賴包

需要添加相應(yīng)的依賴，包括os，torch工具包，torch數(shù)據(jù)載入等依賴

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定義模型案例如下，定義好了方便后續(xù)的調(diào)用

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.定義網(wǎng)絡(luò)架構(gòu)

定義網(wǎng)絡(luò)模型，自定義模型名字，并繼承l(wèi)ightning.pytorch.LightningModule類，如下代碼

• training_step定義了與nn.Module之間交互

• configure_optimizers為模型定義優(yōu)化器

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.定義訓(xùn)練集

定義DataLoader，這一點(diǎn)跟PyTorch調(diào)模型的流程一樣，如下調(diào)用了MNIST公開數(shù)據(jù)集

dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
train_loader = DataLoader(dataset)

5.訓(xùn)練數(shù)據(jù)

使用Lightning來(lái)處理所有的訓(xùn)練，如下代碼。

# model 模型
autoencoder = LitAutoEncoder(Encoder(), Decoder())

# train model 訓(xùn)練
trainer = pl.Trainer()
trainer.fit(model=autoencoder, train_dataloaders=train_loader)

一般的訓(xùn)練過程，需要設(shè)計(jì)如下代碼，進(jìn)行遍歷和循環(huán)訓(xùn)練，Lightning會(huì)消除這些繁瑣的過程，使用Lightning，可以將所有這些技術(shù)混合在一起，而無(wú)需每次都重寫一個(gè)新的循環(huán)。文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-605108.html

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

完整代碼

# 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: 定義一個(gè) 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 定義了一次 預(yù)測(cè)/推理 行為
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        # training_step 定義了一次訓(xùn)練的迭代， 和forward相互獨(dú)立
        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: 定義數(shù)據(jù)集
# -------------------
dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
train, val = data.random_split(dataset, [55000, 5000])

# -------------------
# Step 3: 開始訓(xùn)練
# -------------------
autoencoder = LitAutoEncoder()
trainer = L.Trainer(accelerator="gpu")	
trainer.fit(autoencoder, data.DataLoader(train,batch_size=128), data.DataLoader(val))

到了這里，關(guān)于PyTorch Lightning快速學(xué)習(xí)教程一：快速訓(xùn)練一個(gè)基礎(chǔ)模型的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！