【完全ガイド】PyTorch Apple M1/M2 GPU加速実装 - MPS Backend活用法

# ResNet-50の1エポック訓練時間（ImageNet、バッチサイズ64）
M1 CPU: 1800秒（30分） → 並列度が低い
M1 MPS: 350秒（5.8分） → 5.1倍高速化 ✅
M2 MPS: 280秒（4.7分） → 6.4倍高速化 ✅
NVIDIA RTX 4090: 120秒（2分） → 15倍高速化（参考値）

✅ 解決策一覧

【推奨】解決策1: PyTorch 2.0以降 + MPS Backend導入

最も確実で効果的な方法です。

ステップ1: 環境要件の確認

# macOS バージョン確認（12.3以降が必須）
sw_vers

# 出力例:
# ProductName: macOS
# ProductVersion: 13.2.1

# M1/M2/M3チップの確認
system_profiler SPHardwareDataType | grep "Chip"

# 出力例:
# Chip: Apple M1 Pro

ステップ2: PyTorch 2.0以降のインストール

# 既存のPyTorchをアンインストール（必要な場合）
pip uninstall torch torchvision torchaudio

# 最新版PyTorchをインストール（CPU版ではなく、MPS対応版）
pip install torch torchvision torchaudio

# インストール確認
python -c "import torch; print(f'PyTorch: {torch.__version__}'); print(f'MPS available: {torch.backends.mps.is_available()}')"

# 正しい出力例:
# PyTorch: 2.2.1
# MPS available: True ← これが重要！

ステップ3: 基本的なMPS使用方法

import torch
import torch.nn as nn

# デバイス選択（MPS > CUDA > CPU の優先順位）
if torch.backends.mps.is_available():
    device = torch.device("mps")
    print("✅ MPS device available")
elif torch.cuda.is_available():
    device = torch.device("cuda")
    print("✅ CUDA device available")
else:
    device = torch.device("cpu")
    print("⚠️ CPU only")

# モデルとデータをMPSに転送
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
).to(device)

# 訓練ループ
for epoch in range(10):
    data = torch.randn(64, 784).to(device)
    target = torch.randint(0, 10, (64,)).to(device)
    
    output = model(data)
    loss = nn.CrossEntropyLoss()(output, target)
    
    loss.backward()  # ← MPSで自動的に高速化される！

【性能向上】解決策2: パフォーマンス最適化テクニック

1. 混合精度訓練の活用

FP32（32ビット浮動小数点）からFP16（16ビット）への部分的な変換で、メモリ使用量を半減 + 高速化。

from torch.cuda.amp import autocast, GradScaler

# MPSでもautocastが動作（PyTorch 2.1以降）
scaler = GradScaler()

for epoch in range(epochs):
    for data, target in dataloader:
        data, target = data.to(device), target.to(device)
        
        optimizer.zero_grad()
        
        # 混合精度で forward pass
        with autocast(device_type='mps', dtype=torch.float16):
            output = model(data)
            loss = criterion(output, target)
        
        # スケール調整しながら backward
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

# 効果: ResNet-50で約1.3倍高速化 + メモリ使用量40%削減

2. メモリ効率化テクニック

# グラディエントチェックポイントでメモリ削減
from torch.utils.checkpoint import checkpoint

class OptimizedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(1024, 1024)
        self.layer2 = nn.Linear(1024, 1024)
        self.layer3 = nn.Linear(1024, 10)
    
    def forward(self, x):
        # 中間層でチェックポイント利用（メモリ節約）
        x = checkpoint(self._layer1_forward, x, use_reentrant=False)
        x = checkpoint(self._layer2_forward, x, use_reentrant=False)
        return self.layer3(x)
    
    def _layer1_forward(self, x):
        return torch.relu(self.layer1(x))
    
    def _layer2_forward(self, x):
        return torch.relu(self.layer2(x))

# バッチサイズの動的調整
BATCH_SIZE = 128 if device.type == 'mps' else 64  # MPSでは大きめに設定可能

3. DataLoaderの最適化

from torch.utils.data import DataLoader

# MPSに最適化されたDataLoader設定
dataloader = DataLoader(
    dataset,
    batch_size=128,
    num_workers=4,        # M1/M2では4-8が最適
    pin_memory=False,     # MPSではFalseに設定（重要！）
    persistent_workers=True,  # ワーカープロセス再利用
    prefetch_factor=2     # 先読みバッファ
)

# ❌ 間違い: pin_memory=True（CUDA専用、MPSでは逆効果）
# ✅ 正しい: pin_memory=False（MPS最適化）

【トラブルシューティング】解決策3: 互換性問題の対処

未サポートオペレーションのフォールバック

一部のオペレーションはMPSで未実装の場合があります。その場合はCPUに自動フォールバック。

# 安全なMPS実行ラッパー
def safe_mps_operation(input_tensor, model):
    try:
        # MPSで実行を試みる
        return model(input_tensor.to('mps'))
    except RuntimeError as e:
        if "MPS" in str(e) or "not implemented" in str(e):
            print(f"⚠️ Fallback to CPU: {e}")
            # CPUで実行（遅いが動作する）
            return model(input_tensor.to('cpu'))
        else:
            raise e

# サポート確認スクリプト
def check_mps_support():
    ops_to_test = {
        "Conv2D": lambda: torch.nn.functional.conv2d(
            torch.randn(1, 3, 224, 224, device='mps'),
            torch.randn(64, 3, 7, 7, device='mps'),
            padding=3
        ),
        "BatchNorm2D": lambda: torch.nn.BatchNorm2d(64).to('mps')(
            torch.randn(1, 64, 56, 56, device='mps')
        ),
        "Sparse Operations": lambda: torch.sparse.mm(
            torch.sparse_coo_tensor([[0, 1]], [1.0, 2.0], (2, 2), device='mps'),
            torch.randn(2, 2, device='mps')
        )
    }
    
    for op_name, op_func in ops_to_test.items():
        try:
            op_func()
            print(f"✅ {op_name}: Supported")
        except Exception as e:
            print(f"❌ {op_name}: Not supported - {str(e)[:50]}")

check_mps_support()

既知の制限事項（PyTorch 2.2時点）

オペレーション	MPS対応状況	代替策
Sparse Tensor操作	❌ 一部未実装	CPUフォールバック
カスタムCUDAカーネル	❌ Metal移植が必要	純粋PyTorch OPに書き換え
一部のRNNバリアント	⚠️ 性能低下あり	LSTMは問題なし、GRUは要検証
Conv2D, BatchNorm	✅ 完全サポート	-
Transformer, Attention	✅ 完全サポート	-

【移行ガイド】解決策4: 既存プロジェクトのMPS移行

ステップ1: デバイス抽象化レイヤー作成

# device_utils.py（プロジェクト全体で使い回す）
import torch

def get_best_device():
    """環境に応じた最適なデバイスを自動選択"""
    if torch.backends.mps.is_available():
        if not torch.backends.mps.is_built():
            print("⚠️ MPS not available (PyTorch not built with MPS support)")
            return torch.device("cpu")
        return torch.device("mps")
    elif torch.cuda.is_available():
        return torch.device("cuda")
    else:
        return torch.device("cpu")

# グローバル変数として利用
DEVICE = get_best_device()
print(f"🚀 Using device: {DEVICE}")

ステップ2: 既存CUDAコードの置換

# 変更前（CUDA専用コード）
model = MyModel().cuda()
data = data.cuda()

# 変更後（MPS/CUDA/CPU対応コード）
from device_utils import DEVICE

model = MyModel().to(DEVICE)
data = data.to(DEVICE)

# または、より明示的に
device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
data = data.to(device)

ステップ3: 性能プロファイリング

import time
import torch

def profile_inference(model, sample_data, device_name, num_runs=100):
    """推論速度のベンチマーク"""
    device = torch.device(device_name)
    model = model.to(device)
    sample_data = sample_data.to(device)
    
    # ウォームアップ
    for _ in range(10):
        with torch.no_grad():
            _ = model(sample_data)
    
    # 計測開始
    start = time.time()
    with torch.no_grad():
        for _ in range(num_runs):
            _ = model(sample_data)
    
    # MPSの場合は同期が必要
    if device_name == "mps":
        torch.mps.synchronize()
    
    elapsed = time.time() - start
    avg_time = (elapsed / num_runs) * 1000  # ミリ秒
    
    print(f"Device: {device_name}")
    print(f"  Total: {elapsed:.2f}秒")
    print(f"  Average: {avg_time:.2f}ms")
    print(f"  Throughput: {num_runs / elapsed:.2f} inferences/sec")
    
    return elapsed

# 実行例
from torchvision.models import resnet50

model = resnet50(pretrained=False)
sample = torch.randn(1, 3, 224, 224)

print("=== ResNet-50 Inference Benchmark ===")
cpu_time = profile_inference(model, sample, "cpu")
mps_time = profile_inference(model, sample, "mps")

print(f"\n⚡ Speedup: {cpu_time / mps_time:.2f}x faster with MPS")