pytorch-research

PyTorch - 高度な研究およびエンジニアリング

研究グレードの PyTorch を使うには、nn.Sequential の枠を超える必要があります。勾配の流れや重みの初期化方法、複数の GPU にまたがる計算の分散を制御しなければなりません。このガイドでは、フレームワークの「内部」を解説します。

使用する場面

• 標準とは異なる数学的導関数を持つカスタムレイヤーを実装する場合。
• フックを使って勾配の消失や爆発をデバッグする場合。
• モデルを複数の GPU にスケールする場合（分散データ並列）。
• PyTorch Profiler を使ってモデルのパフォーマンスを微調整する場合。
• 複雑な学習率スケジュール（Cyclic、OneCycle）を作成する場合。
• 高性能な推論向けにモデルをデプロイする場合（TorchScript、FX）。
• 重みの初期化や正規化の手法を研究する場合。

リファレンスドキュメント

• Autograd の仕組み: https://pytorch.org/docs/stable/notes/autograd.html
• 分散トレーニング: https://pytorch.org/docs/stable/distributed.html
• プロファイラー: https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html
• 検索パターン: torch.autograd.Function、register_forward_hook、DistributedDataParallel、torch.nn.init

基本原則

計算グラフを超えて

PyTorch は「define-by-run」方式のフレームワークですが、研究では、グラフを壊さずにバックワードパスへ介入したり、中間テンソルを調べたりする必要があることがよくあります。

勾配のライフサイクル

勾配は .grad 属性に累積され、retain_graph=True を指定しない限り、backward() によってグラフが消費されることを理解しておきましょう。

メモリと速度

研究では、チェックポイント化などの手法を使い、メモリ（アクティベーション）と引き換えに速度（再計算）を得ることがよくあります。

クイックリファレンス

標準インポート

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

基本パターン - カスタム Autograd 関数

class MySignFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # Save input for backward pass
        ctx.save_for_backward(input)
        return torch.sign(input)

    @staticmethod
    def backward(ctx, grad_output):
        # Straight-through estimator (STE) logic
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        # Custom logic: gradients pass through as if it were an identity
        return grad_input

# Usage
my_sign = MySignFunction.apply

重要なルール

✅ 推奨事項

• register_full_backward_hook を使用する - 特定のモジュールを通る勾配を調べたり変更したりできます。
• 重みを明示的に初期化する - model.apply(fn) ループ内で torch.nn.init（Xavier、Kaiming）を使用します。
• DistributedDataParallel（DDP）を使用する - DataParallel（DP）の代わりに使用します。DDP はより高速で、マルチプロセス間の通信を適切に処理します。
• 最適化の前にプロファイルする - torch.profiler を使って、実際にモデルを遅くしている演算子（遅い view() や cat() など）を特定します。
• torch.cuda.empty_cache() の使用は控える - OS に物理メモリを解放するわけではなく、PyTorch のメモリマネージャーを断片化させます。必要な場合に限り、長時間実行するループ内で使用してください。

❌ 非推奨事項

• カスタムレイヤーで inplace=True を使用しない - Autograd が勾配を正しく計算できなくなることがよくあります。
• ループ内で item() を使用しない - GPU テンソルに対して .item() を呼び出すと、CPU と GPU の同期が強制され、パフォーマンスが大幅に低下します。
• DistributedSampler で shuffle=False を設定し忘れない - マルチ GPU 環境では、シャッフル処理をサンプラーに任せます。
• グローバル変数を避ける - PyTorch モデルは、シリアライズやデプロイを簡単にするため、自己完結型にする必要があります。

高度なカスタムレイヤー

デバッグと特徴抽出のためのフック

def print_grad_norm(module, grad_input, grad_output):
    print(f"Module: {module.__class__.__name__}, Grad Norm: {grad_output[0].norm().item()}")

# Attach to a specific layer
model.fc1.register_full_backward_hook(print_grad_norm)

# Extract activations (Forward Hook)
activations = {}
def get_activation(name):
    def hook(model, input, output):
        activations[name] = output.detach()
    return hook

model.conv1.register_forward_hook(get_activation('conv1'))

高度なトレーニングパターン

分散データ並列（DDP）の基本構成

import torch.multiprocessing as mp

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def train(rank, world_size):
    setup(rank, world_size)
    
    model = MyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    # Use DistributedSampler to ensure each GPU sees different data
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = DataLoader(dataset, sampler=sampler, batch_size=32)
    
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=0.001)
    # ... training loop ...
    
    dist.destroy_process_group()

# mp.spawn(train, args=(world_size,), nprocs=world_size)

パフォーマンスとプロファイリング

プロファイラーの使用

from torch.profiler import profile, record_function, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], 
             record_shapes=True) as prof:
    with record_function("model_inference"):
        model(inputs)

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

勾配チェックポイント化（メモリ節約）

非常に深いモデルを扱っていてメモリが限られている場合は、計算と引き換えにメモリを節約します。

from torch.utils.checkpoint import checkpoint

class DeepModel(nn.Module):
    def forward(self, x):
        # Instead of storing all activations, recompute them during backward
        x = checkpoint(self.heavy_layer_1, x)
        x = checkpoint(self.heavy_layer_2, x)
        return x

実践的なワークフロー

1. カスタム重み初期化

def init_weights(m):
    if isinstance(m, nn.Linear):
        # Kaiming initialization for ReLU networks
        torch.nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        if m.bias is not None:
            torch.nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.Conv2d):
        torch.nn.init.xavier_uniform_(m.weight)

model.apply(init_weights)

2. 勾配クリッピング（安定性）

# Inside training loop
loss.backward()

# Clip to prevent exploding gradients (standard in RNNs/Transformers)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

optimizer.step()

3. 動的学習率（OneCycleLR）

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.1, 
                                                steps_per_epoch=len(train_loader), 
                                                epochs=10)

for epoch in range(10):
    for batch in train_loader:
        train_batch()
        scheduler.step() # Step every batch for OneCycle

よくある落とし穴と解決策

インプレース変更エラー

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation.

# ❌ Problem: x += 1 (breaks backward pass)
# ✅ Solution: y = x + 1 (creates a new tensor)

CUDA のメモリ不足（OOM）対策

• バッチサイズ：小さくします。
• 勾配蓄積：小さなバッチで損失を計算し、N ステップごとにのみ step() を実行します。
• キャッシュのクリア：独立した評価の間に torch.cuda.empty_cache() を使用します。
• 混合精度：torch.cuda.amp を使用します（メモリを 50% 節約できます）。

サイレントな失敗：zero_grad() の位置

backward() の後、step() の前に zero_grad() を呼び出すと、モデルは学習しません。

# ✅ Correct order:
optimizer.zero_grad()
loss.backward()
optimizer.step()

Research PyTorch では、数学的エンジンを深く理解することが重要です。カスタム勾配、フック、分散インフラストラクチャを活用することで、標準的なモデルのトレーニングから、次世代の科学的 AI の創出へと進むことができます。