15分後に雷が落ちる

15 分で Lightning — PyTorch Lightning 2.0.4 ドキュメント

PyTorch Lightningをインストールする

pip install lightning

conda install lightning -c conda-forge

LightningModuleを定義する

LightningModuleいくつかのトレーニングプロセス (および検証とテスト)pytorchを統合することができます。nn.Module

以下は、手書き数字認識用のオートエンコーダーの例です。

import os
import torch
from torch import optim, nn, utils
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
import lightning.pytorch as pl

'''
定义两个模型，编码器和解码器，这个是pytorch的模型对象
'''
encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))

# 定义LightningModule
class LitAutoEncoder(pl.LightningModule):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def training_step(self, batch, batch_idx):
        # 训练步骤
        # 这个跟 forward 不相关
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = nn.functional.mse_loss(x_hat, x)
        # 存储日志（需要安装Tensorboard）
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
				# 优化器
        optimizer = optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

# 初始化自动编码器
autoencoder = LitAutoEncoder(encoder, decoder)

データセットを定義する

Lightningすべての反復可能なデータセット形式 ( DataLoader、numpy、その他) がサポートされています。

# setup 
datadataset=MNIST(os.getcwd(),download=True,transform=ToTensor())
train_loader=utils.data.DataLoader(dataset)

トレーニングモデル

LightningオブジェクトはさまざまなデータセットとTrainer統合できLightningModule、拡張にはいくつかのエンジニアリング手法が必要です。

# 训练模型
trainer = pl.Trainer(limit_train_batches=100, max_epochs=10)
trainer.fit(model=autoencoder, train_dataloaders=train_loader)

Trainerこのオブジェクトには、一般的に使用される多くのプロシージャも実装されています。

Epochそしてbatchそれを繰り返します。
optimizer.step()、、 loss.backward()_ optimizer.zero_grad()
** 確認中ですmodel.eval()。**
モデルの保存とロード
テンソルボード
マルチGPU
TPU
半精度ブレンド

python[注]: jupyter では、マルチカードトレーニングでエラーが報告される場合があります。コードを直接使用してみてください。

使用モデル

モデルをトレーニングした後、onnx や torchscript にエクスポートして本番環境に導入することも、重みをロードして予測を実行することもできます。

# 载入模型
checkpoint = "./lightning_logs/version_0/checkpoints/epoch=0-step=100.ckpt"
autoencoder = LitAutoEncoder.load_from_checkpoint(checkpoint, encoder=encoder, decoder=decoder)

# 选择训练好的编码器
encoder = autoencoder.encoder
encoder.eval()

# 编码图片
fake_image_batch = torch.randn(8, 28 * 28).to(next(encoder.parameters()).device)
embeddings = encoder(fake_image_batch)
print("⚡" * 20, "\nPredictions (4 image embeddings):\n", embeddings, "\n", "⚡" * 20)

トレーニングの視覚化

インストールされている場合はTensorboard、実験プロセスを観察するために使用できます。

tensorboard --logdir .

追加のトレーニング設定

# 4gpu训练
trainer = Trainer(
    devices=4,
    accelerator="gpu",
 )

# train 1TB+ parameter models with Deepspeed/fsdp
# 使用 Deepspeed 训练大模型
trainer = Trainer(
    devices=4,
    accelerator="gpu",
    strategy="deepspeed_stage_2",
    precision=16
 )

# 20+ helpful flags for rapid idea iteration
# 有助于快速迭代的一些设置
trainer = Trainer(
    max_epochs=10,
    min_epochs=5,
    overfit_batches=1
 )

# access the latest state of the art techniques
# 获取最新的技术
trainer = Trainer(callbacks=[StochasticWeightAveraging(...)])

いくつかの柔軟な設定

カスタムトレーニングループ

LightningModuleには 20 種類を超えるブレークポイント ( HOOK )が設定されており、これを使用してトレーニングプロセスをカスタマイズできます。

class LitAutoEncoder(pl.LightningModule):
    def backward(self, loss):
        loss.backward()

エクステンドトレーナー

上記のコードでは、モデルのストレージ用にいくつかの設定が行われています。これらの設定はpl.Callbackオブジェクトに実装したり、オブジェクトにインポートしたりできますTrainer。

次のようにインポートできますTrainer。

trainer = Trainer(callbacks=[AWSCheckpoints()])

PyTorch Lightning の基本を開始する