PyTorch入門 第8回：モデル保存、ファインチューニング、GPUとよくあるエラー

第7回：訓練を安定させる
  -> Weight Decayで重みを制御する
  -> Dropoutで活性値へランダム性を入れる
  -> Normalizationで特徴の尺度を整える

第8回：訓練結果を運用できる形にする
  -> state_dictとして保存する
  -> checkpointから訓練を再開する
  -> 事前学習済みモデルを別タスクへ使う
  -> GPUや分散環境で安全に動かす
  -> エラーを体系的に切り分ける

訓練中に管理したいもの
├── model parameters / buffers
├── optimizerの内部状態
├── schedulerの進み具合
├── AMPのGradScaler状態
├── epoch / global_step
├── best metric
├── class_to_idxや前処理設定
└── 乱数状態、設定、環境メモ

保存する
  model / optimizer / scheduler / scaler / metadata
      |
      v
復元する
  オブジェクトを作り直す -> state_dictを読み込む -> deviceへ移す
      |
      v
使う
  推論 / 評価 / 追加訓練 / ファインチューニング

メモリ上のTensorや辞書
  -- torch.save -->
ディスク上の .pt / .pth ファイル
  -- torch.load -->
PythonオブジェクトやTensor状態

torch.save(obj, f)

state = torch.load(
    path,
    map_location="cpu",
    weights_only=True,
)

torch.save(model.state_dict(), "model_weights.pt")

model = MyModel(...)
state = torch.load("model_weights.pt", map_location="cpu", weights_only=True)
model.load_state_dict(state)
model.eval()

from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from torch import nn


class TinyClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(4, 3),  # 4次元入力を3次元へ
            nn.ReLU(),
            nn.Linear(3, 2),  # 3次元から2クラス分のlogitへ
        )

    def forward(self, x):
        return self.net(x)


torch.manual_seed(0)

with TemporaryDirectory() as workdir:
    path = Path(workdir) / "tiny_weights.pt"

    model = TinyClassifier()
    torch.save(model.state_dict(), path)  # 重みだけをファイルへ書き出す

    state = torch.load(path, map_location="cpu", weights_only=True)
    restored = TinyClassifier()  # modelとは別の、新しいインスタンス
    missing_keys, unexpected_keys = restored.load_state_dict(state)  # 読み込んだ重みをrestoredへ流し込む
    restored.eval()

    print("keys:", list(state.keys()))
    # keys: ['net.0.weight', 'net.0.bias', 'net.2.weight', 'net.2.bias']
    print("missing:", missing_keys, "unexpected:", unexpected_keys)
    # missing: [] unexpected: []

from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from torch import nn


class LeNetLike(nn.Module):
    def __init__(self, classes):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
        )  # 畳み込みとpoolingで特徴マップを抽出する部分
        self.classifier = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, classes),
        )  # 抽出した特徴からクラスを判定する部分

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)  # batch次元を残したままフラット化する
        return self.classifier(x)

    @torch.no_grad()
    def initialize_constant(self, value):
        # デモ専用: 全パラメータを同じ定数にして「訓練後の重み」を再現する（実際の訓練では使わない）
        for parameter in self.parameters():
            parameter.fill_(value)


torch.manual_seed(0)

with TemporaryDirectory() as workdir:
    workdir = Path(workdir)
    state_path = workdir / "model_state_dict.pkl"

    model = LeNetLike(classes=2)
    print("before:", round(model.features[0].weight[0, 0, 0, 0].item(), 4))
    # before: -0.0009
    model.initialize_constant(20191104.0)  # 「訓練後の重み」として全パラメータを同じ値にする
    print("after:", round(model.features[0].weight[0, 0, 0, 0].item(), 1))
    # after: 20191104.0

    torch.save(model.state_dict(), state_path) # 推奨しやすい保存形式

    restored = LeNetLike(classes=2)
    restored.initialize_constant(0.0)  # 読み込み前の値をわざと0にして、復元されたことを確認しやすくする
    state = torch.load(state_path, map_location="cpu", weights_only=True)
    restored.load_state_dict(state)

    print("state keys head:", list(state.keys())[:3])
    # state keys head: ['features.0.weight', 'features.0.bias', 'features.3.weight']
    print("restored:", round(restored.features[0].weight[0, 0, 0, 0].item(), 1))
    # restored: 20191104.0
    print("file:", state_path.name)
    # file: model_state_dict.pkl

checkpoint = {
    "format_version": 1,
    "epoch": epoch,
    "global_step": global_step,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "scaler": scaler.state_dict() if scaler is not None else None,
    "best_metric": best_metric,
    "config": config,
    "class_to_idx": class_to_idx,
}
torch.save(checkpoint, "checkpoint.pt")

temporary = checkpoint_path.with_suffix(".tmp")
torch.save(checkpoint, temporary)
temporary.replace(checkpoint_path)

from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


torch.manual_seed(0)

features = torch.randn(32, 4)
targets = (features.sum(dim=1) > 0).long()
loader = DataLoader(TensorDataset(features, targets), batch_size=8, shuffle=False)


def build_objects():
    # 毎回まっさらな状態のmodel/optimizer/schedulerをまとめて作るヘルパー
    model = nn.Sequential(
        nn.Linear(4, 8),
        nn.ReLU(),
        nn.Linear(8, 2),
    )
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)
    return model, optimizer, scheduler


def train_one_epoch(model, optimizer, criterion):
    # loaderを1周してパラメータを更新し、平均lossを返す
    model.train()
    total_loss = 0.0
    for inputs, labels in loader:
        logits = model(inputs)
        loss = criterion(logits, labels)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)


criterion = nn.CrossEntropyLoss()

with TemporaryDirectory() as workdir:
    checkpoint_path = Path(workdir) / "checkpoint_2_epoch.pt"
    model, optimizer, scheduler = build_objects()

    for epoch in range(3):
        loss = train_one_epoch(model, optimizer, criterion)
        scheduler.step()
        if epoch == 2:  # epoch 2まで進めた時点のcheckpointを保存する
            torch.save(
                {
                    "epoch": epoch,
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "scheduler_state_dict": scheduler.state_dict(),
                    "loss": loss,
                },
                checkpoint_path,
            )

    restored_model, restored_optimizer, restored_scheduler = build_objects()  # 新しいオブジェクト群を作り、訓練継続をゼロから再現する
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    restored_model.load_state_dict(checkpoint["model_state_dict"])
    restored_optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    restored_scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
    start_epoch = checkpoint["epoch"]

    print("resume from epoch:", start_epoch)
    # resume from epoch: 2
    print("restored lr:", restored_optimizer.param_groups[0]["lr"])
    # restored lr: 0.010000000000000002

    for epoch in range(start_epoch + 1, 5):
        loss = train_one_epoch(restored_model, restored_optimizer, criterion)
        restored_scheduler.step()
        print("epoch:", epoch, "loss:", round(loss, 4), "lr:", restored_optimizer.param_groups[0]["lr"])
        # epoch: 3 loss: 0.4695 lr: 0.0010000000000000002
        # epoch: 4 loss: 0.4569 lr: 0.0010000000000000002

state = torch.load(
    "model_weights.pt",
    map_location="cpu",
    weights_only=True,
)

state = torch.load(
    "model_weights.pt",
    map_location="cpu",
    weights_only=True,
)

state = torch.load(
    "checkpoint.pt",
    map_location={"cuda:1": "cuda:0"},
    weights_only=True,
)

import torch


large = torch.arange(1_000_000)
small_view = large[:5]
small_to_save = small_view.clone()  # 必要な5要素だけを独立したTensorにする

result = model.load_state_dict(state, strict=False)
print(result.missing_keys)
print(result.unexpected_keys)

import torch
from torch import nn


class TinyTransferModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = nn.Linear(4, 3)  # 事前学習済みとして扱う部分
        self.head = nn.Linear(3, 2)  # タスクに合わせて後から差し替える部分

    def forward(self, x):
        x = torch.relu(self.backbone(x))
        return self.head(x)


model = TinyTransferModel()
state = model.state_dict()
backbone_only_state = {
    key: value
    for key, value in state.items()
    if not key.startswith("head.")
}  # "head."で始まらないキーだけを残す(=headの重みを除く)

result = model.load_state_dict(backbone_only_state, strict=False)
print("missing:", result.missing_keys)
# missing: ['head.weight', 'head.bias']
print("unexpected:", result.unexpected_keys)
# unexpected: []

raw_model = model.module if hasattr(model, "module") else model
torch.save(raw_model.state_dict(), "model_weights.pt")

import torch


state = {
    "module.backbone.weight": torch.ones(2, 2),
    "module.head.bias": torch.zeros(2),
}  # DataParallel/DDPで保存された想定のキー
cleaned_state = {
    key.removeprefix("module."): value
    for key, value in state.items()
}  # 先頭の"module."だけを取り除く

print(list(cleaned_state.keys()))
# ['backbone.weight', 'head.bias']

モデルを作る
 -> optimizer / scheduler / scalerを作る
 -> checkpointを読み込む
 -> model.load_state_dict
 -> optimizer.load_state_dict
 -> scheduler / scalerを復元する
 -> epochやglobal_stepを復元する
 -> 必要なら乱数状態も復元する
 -> 訓練を続ける

rng_state = {
    "torch": torch.get_rng_state(),
    "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
}

訓練checkpoint:
  訓練状態をディスクへ保存し、後で復元する

activation checkpointing:
  順伝播の中間活性を保存しすぎず、逆伝播時に再計算してGPUメモリを節約する

事前学習済みモデル
├── backbone / feature extractor
└── task head / classifier

from torchvision.models import resnet18, ResNet18_Weights


weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)
preprocess = weights.transforms()

import torch
from torch import nn


class TinyFineTuneModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Linear(4, 4),
            nn.ReLU(),
        )  # 事前学習済みとして扱う部分(ここでは凍結する)
        self.head = nn.Linear(4, 2)  # タスク用に訓練する部分

    def forward(self, x):
        return self.head(self.backbone(x))


model = TinyFineTuneModel()
for parameter in model.backbone.parameters():
    parameter.requires_grad_(False)  # backboneの勾配計算を止める

optimizer = torch.optim.AdamW(
    (parameter for parameter in model.parameters() if parameter.requires_grad),  # 勾配計算が有効なパラメータだけをoptimizerに渡す
    lr=1e-3,
)

trainable_names = [
    name for name, parameter in model.named_parameters()
    if parameter.requires_grad
]
print(trainable_names)
# ['head.weight', 'head.bias']

optimizer = torch.optim.AdamW(
    [
        {"params": model.layer4.parameters(), "lr": 1e-4},  # backbone終盤のblockは小さめの学習率
        {"params": model.fc.parameters(), "lr": 1e-3},       # 新しく置き換えたheadは大きめの学習率
    ],
    weight_decay=1e-2,
)

from torch import nn


model.train()  # モデル全体は訓練モードにする（Dropoutなどは有効なまま）
for module in model.backbone.modules():
    if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
        module.eval()  # BatchNormだけrunning statisticsの更新を止める

import torch
from torch import nn
from torchvision.models import resnet18, ResNet18_Weights


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

weights = ResNet18_Weights.DEFAULT
model = resnet18(weights=weights)  # ImageNet事前学習済みの重みを読み込む

for parameter in model.parameters():
    parameter.requires_grad_(False)  # まずは全パラメータの勾配計算を止める

in_features = model.fc.in_features
model.fc = nn.Linear(in_features, 2)  # 新しいLinearに置き換える(新しいパラメータはrequires_grad=Trueが既定)
model.to(device)

optimizer = torch.optim.AdamW(model.fc.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
preprocess = weights.transforms()

import torch
from torch import nn
from torchvision.models import resnet18


torch.manual_seed(0)
classes = 2
lr = 1e-3
model = resnet18(weights=None)  # 構造確認用。実データでは事前学習済みweightsを指定する

num_features = model.fc.in_features
model.fc = nn.Linear(num_features, classes)  # 最後の分類headを2クラス用に差し替える

fc_param_ids = {id(parameter) for parameter in model.fc.parameters()}
base_params = [
    parameter for parameter in model.parameters()
    if id(parameter) not in fc_param_ids
]
optimizer = torch.optim.SGD(
    [
        {"params": base_params, "lr": lr * 0.1},      # backboneは小さめの学習率
        {"params": model.fc.parameters(), "lr": lr}, # 新しいheadは大きめの学習率
    ],
    momentum=0.9,
)

example = torch.randn(2, 3, 224, 224)
logits = model(example)
print("fc in_features:", num_features)
# fc in_features: 512
print("fc out_features:", model.fc.out_features)
# fc out_features: 2
print("param group lrs:", [group["lr"] for group in optimizer.param_groups])
# param group lrs: [0.0001, 0.001]
print("logits shape:", logits.shape)
# logits shape: torch.Size([2, 2])

import torch


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

import torch


print("cuda available:", torch.cuda.is_available())
# cuda available: False
print("cuda device count:", torch.cuda.device_count())
# cuda device count: 0
if torch.cuda.is_available():
    print("device 0:", torch.cuda.get_device_name(0))

import torch


gpu_id = 0
gpu_str = f"cuda:{gpu_id}"
device = torch.device(gpu_str if torch.cuda.is_available() else "cpu")

x_cpu = torch.ones((3, 3))
x_device = x_cpu.to(device)

print("device:", x_device.device, "is_cuda:", x_device.is_cuda, "shape:", tuple(x_device.shape))
# device: cpu is_cuda: False shape: (3, 3)
print("device_count:", torch.cuda.device_count())
# device_count: 0
if torch.cuda.is_available():
    print("device_name:", torch.cuda.get_device_name(0))
else:
    print("device_name: CPU only")
    # device_name: CPU only

import torch
from torch import nn


x = torch.ones(2)
y = x.to(dtype=torch.float64)

print(x.dtype, y.dtype, x is y)
# torch.float32 torch.float64 False

linear = nn.Linear(2, 1)
returned = linear.to("cpu")
print(returned is linear)
# True

model = MyModel().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

mask = images.new_zeros(images.shape[0], dtype=torch.bool)

images = images.to(device, non_blocking=True)
targets = targets.to(device, non_blocking=True)

CUDA_VISIBLE_DEVICES=2,3 python train.py

物理GPU 2 -> 論理 cuda:0
物理GPU 3 -> 論理 cuda:1

torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)

if torch.cuda.is_available():
    allocated = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    print("allocated MB:", round(allocated, 2))
    print("reserved MB:", round(reserved, 2))

scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda")  # CUDA以外ではno-opになる

for images, targets in train_loader:
    images = images.to(device, non_blocking=True)
    targets = targets.to(device, non_blocking=True)
    optimizer.zero_grad(set_to_none=True)

    with torch.autocast(
        device_type=device.type,
        dtype=torch.float16,
        enabled=device.type == "cuda",
    ):  # このwithの中だけ、対象演算がfloat16で実行される
        logits = model(images)
        loss = criterion(logits, targets)

    scaler.scale(loss).backward()  # lossをscaleしてからbackwardする
    scaler.step(optimizer)         # 勾配をunscaleしてからoptimizer.step相当を行う
    scaler.update()                # 次iterationのscaleを調整する

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()

model = nn.DataParallel(model)

import os
import torch
from torch import nn


os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class FooNet(nn.Module):
    def __init__(self, neural_num, layers=3):
        super().__init__()
        self.linears = nn.ModuleList([
            nn.Linear(neural_num, neural_num, bias=False)
            for _ in range(layers)
        ])

    def forward(self, x):
        print("batch size in forward:", x.size(0))
        # batch size in forward: 16
        for linear in self.linears:
            x = torch.relu(linear(x))
        return x


inputs = torch.randn(16, 3).to(device)
model = FooNet(neural_num=3, layers=3)
model = nn.DataParallel(model).to(device)
outputs = model(inputs)

print("model outputs.size:", outputs.size())
# model outputs.size: torch.Size([16, 3])
print("CUDA_VISIBLE_DEVICES:", os.environ["CUDA_VISIBLE_DEVICES"])
# CUDA_VISIBLE_DEVICES: 0
print("device_count:", torch.cuda.device_count())
# device_count: 0

主GPU上の入力batch
 -> batch次元でscatter
 -> 各GPUへmodel replicaを作る
 -> forward / backward
 -> 勾配を主モデルへ集約
 -> 出力をgather

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP


dist.init_process_group(device_id=local_rank)  # プロセス間の通信グループを初期化する
torch.cuda.set_device(local_rank)  # このプロセスが使うGPUを指定する

model = MyModel().to(local_rank)
model = DDP(model, device_ids=[local_rank])  # 勾配のall-reduceを自動化するラッパー

torchrun --standalone --nnodes=1 --nproc-per-node=gpu train.py

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)  # rankごとにデータを分担させる
train_loader = DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size)

for epoch in range(num_epochs):
    train_sampler.set_epoch(epoch)  # epochごとにシャッフル順を変える(再現性のため)
    train_one_epoch(...)

DataParallel:
1プロセス
batch -> scatter -> GPU replica -> gather -> 主GPUでloss

DDP:
rank 0: sampler shard 0 -> GPU 0 -> local loss -> gradients --+
rank 1: sampler shard 1 -> GPU 1 -> local loss -> gradients --+-> all-reduce
rank 2: sampler shard 2 -> GPU 2 -> local loss -> gradients --+

import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset, TensorDataset
import torchvision.transforms as transforms


try:
    DataLoader(
        TensorDataset(torch.empty(0, 3), torch.empty(0, dtype=torch.long)),
        batch_size=4,
        shuffle=True,
    )
except ValueError as error:
    print("empty dataset:", str(error).splitlines()[0])
    # empty dataset: num_samples should be a positive integer value, but got num_samples=0

try:
    transforms.ToTensor()(torch.zeros(3, 32, 32))
except TypeError as error:
    print("totensor tensor:", str(error).splitlines()[0])
    # totensor tensor: pic should be PIL Image or ndarray. Got <class 'torch.Tensor'>


class VariableSizeDataset(Dataset):
    def __len__(self):
        return 2

    def __getitem__(self, index):
        size = 32 + index
        return torch.zeros(3, size, size), torch.tensor(index)


try:
    next(iter(DataLoader(VariableSizeDataset(), batch_size=2)))
except RuntimeError as error:
    print("collate shape:", str(error).splitlines()[0])
    # collate shape: stack expects each tensor to be equal size, but got [3, 32, 32] at entry 0 and [3, 33, 33] at entry 1

conv = nn.Conv2d(3, 6, kernel_size=5)
try:
    conv(torch.randn(4, 1, 32, 32))
except RuntimeError as error:
    print("conv channel:", str(error).splitlines()[0])
    # conv channel: Given groups=1, weight of size [6, 3, 5, 5], expected input[4, 1, 32, 32] to have 3 channels, but got 1 channels instead


class FooNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(3, 3)

    def forward(self, x):
        return self.linear(x)


wrapped = nn.DataParallel(FooNet())
try:
    wrapped.linear
except AttributeError as error:
    print("dataparallel attr:", str(error).splitlines()[0])
    # dataparallel attr: 'DataParallel' object has no attribute 'linear'

criterion = nn.CrossEntropyLoss()
try:
    criterion(torch.randn(3, 2), torch.tensor([0, 1, 2]))
except (IndexError, RuntimeError) as error:
    print("target range:", str(error).splitlines()[0])
    # target range: Target 2 is out of bounds.

ValueError: num_samples should be a positive integer value,
but got num_samples=0

print(len(dataset))
print(getattr(dataset, "samples", [])[:5])

TypeError: pic should be PIL Image or ndarray. Got torch.Tensor

Sizes of tensors must match...

print(x.shape)

AttributeError: 'DataParallel' object has no attribute 'linear'

Attempting to deserialize object on a CUDA device...

state = torch.load(path, map_location="cpu", weights_only=True)

Can't get attribute 'FooNet'

target >= 0 && target < n_classes

Expected all tensors to be on the same device
Expected scalar type ...

DataLoader worker ... killed by signal: Killed

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

torch.autograd.set_detect_anomaly(True)

assert torch.isfinite(loss)

1. traceback全体を読む
2. num_workers=0にする
3. 単GPUまたはCPUで再現する
4. AMPやcompileを一度外す
5. 1 batch、さらに1 sampleへ縮める
6. Dataset -> transform -> collate -> model -> lossの境界で確認する
7. shape、dtype、device、min/max、finiteを出す
8. 公式の最小例と比べる
9. 複雑な機能を1つずつ戻す

assert images.ndim == 4
assert images.shape[1] == 3
assert targets.dtype == torch.long
assert images.device == next(model.parameters()).device
assert torch.isfinite(images).all()

from pathlib import Path
import torch


checkpoint_path = Path("checkpoints/last.pt")
checkpoint_path.parent.mkdir(parents=True, exist_ok=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = build_model().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda")

start_epoch = 0
best_metric = float("-inf")

if checkpoint_path.exists():  # 前回の続きがあれば読み込む
    checkpoint = torch.load(
        checkpoint_path,
        map_location="cpu",
        weights_only=True,
    )
    model.load_state_dict(checkpoint["model"])
    optimizer.load_state_dict(checkpoint["optimizer"])
    scheduler.load_state_dict(checkpoint["scheduler"])
    scaler.load_state_dict(checkpoint["scaler"])
    start_epoch = checkpoint["epoch"] + 1  # 保存時のepochの次から再開する
    best_metric = checkpoint["best_metric"]
    model.to(device)

for epoch in range(start_epoch, num_epochs):
    model.train()
    for images, targets in train_loader:
        images = images.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)
        optimizer.zero_grad(set_to_none=True)

        with torch.autocast(
            device_type=device.type,
            dtype=torch.float16,
            enabled=device.type == "cuda",
        ):
            logits = model(images)
            loss = criterion(logits, targets)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

    scheduler.step()
    metric = evaluate(model, valid_loader)
    best_metric = max(best_metric, metric)

    raw_model = model.module if hasattr(model, "module") else model  # DataParallel/DDPで包まれていれば中身を取り出す
    checkpoint = {
        "format_version": 1,
        "epoch": epoch,
        "model": raw_model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
        "scaler": scaler.state_dict(),
        "best_metric": best_metric,
    }
    temporary = checkpoint_path.with_suffix(".tmp")
    torch.save(checkpoint, temporary)  # まず一時ファイルへ保存してから
    temporary.replace(checkpoint_path)  # 置き換える(保存中断による破損を避ける)

訓練中の状態
├── model parameters / buffers
├── optimizer moments
├── scheduler position
├── AMP scaler
└── epoch / best metric
        |
        v torch.save
checkpoint.pt
        |
        v torch.load + load_state_dict
復元された訓練状態

事前学習済みモデル
  -> 汎用特徴を持つbackbone
  -> 目標タスク用headへ置き換える
  -> freeze / partial unfreeze / full fine-tune

CUDA_VISIBLE_DEVICES=2,3

物理GPU 2 -> 論理 cuda:0
物理GPU 3 -> 論理 cuda:1

batch 16
 -> scatter: 4 + 4 + 4 + 4
 -> 4つのreplica
 -> gather: [16, C]

4 processes x 1 GPU
 -> 各rankが別shardを読む
 -> local forward / backward
 -> gradient all-reduce
 -> 各rankでoptimizer.step()

Dataset -> transform -> collate -> model -> loss -> backward -> optimizer
   |          |           |        |       |        |          |
 path       dtype       shape    device  target   NaN/Inf    state