振り返り: PyTorch Lightning によるスケーラブルな AI モデル

PyTorch Lightning によるスケーラブルな AI モデルの構築

Sergiy Tkachuk

Director, GenAI Productivity

PyTorch Lightning でスケーラブルなモデルを構築する

ニューラルネットワークの構造化
効果的なトレーニングロジック
LightningModule の実装

import lightning.pytorch as pl
import torch.nn as nn

class ClassificationModel(pl.LightningModule):
    def __init__(self, input_dim,
                 hidden_dim, num_class):
          # Initialize parent class
        super().__init__()
        # First layer
        self.layer1 = nn.Linear(input_dim,
                                hidden_dim)
        # Activation function
        self.relu = nn.ReLU()
        # Output layer
        self.layer2 = nn.Linear(hidden_dim,
                                num_class)

PyTorch Lightning の高度なテクニック

データ管理と前処理
検証とテストの手法
Lightning コールバックの最適化

from lightning.pytorch import Trainer
from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint

checkpoint = ModelCheckpoint(
    monitor='val_accuracy',
    save_top_k=2,
    mode='max')
early_stopping = EarlyStopping(
    monitor='val_accuracy',
    patience=5,
    mode='max')

スケーラビリティのためのモデル最適化

モデル量子化の戦略
効率化のためのモデル剪定
デプロイ向け TorchScript

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def forward(self, x):
        return x * 2

scripted_model = torch.jit.script(SimpleModel())

torch.jit.save(scripted_mod,"model.pt") # Save the model
loaded_model=torch.jit.load("model.pt") # Load the model

おめでとうございます！

PyTorch Lightning によるスケーラブルな AI モデルの構築