PyTorch入門 第6回：学習率スケジューラ、TensorBoard、HookとGrad-CAMで訓練を観察する

optimizer.zero_grad(set_to_none=True)
logits = model(inputs)
loss = criterion(logits, targets)
loss.backward()
optimizer.step()

for epoch:
  train loop
    forward / loss / backward / optimizer.step
    TensorBoardへlossやlrを書く
    必要ならbatch単位scheduler.step

  validation
    valid lossやmetricを計算する
    epoch単位scheduler.step、またはReduceLROnPlateauへmetricを渡す
    checkpointへmodel/optimizer/schedulerを保存する

必要なとき:
  Hookで中間特徴や勾配を取り出す
  Grad-CAMで画像上の注目領域を確認する

新しい重み = 今の重み - 学習率 * 勾配

torch.optim.lr_scheduler.LRScheduler(optimizer, last_epoch=-1)

optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
scheduler.step()

valid_loss = evaluate(model, valid_loader)
scheduler.step(valid_loss)

current_lrs = scheduler.get_last_lr()
current_lr = optimizer.param_groups[0]["lr"]

torch.optim.lr_scheduler.StepLR(
    optimizer,
    step_size,
    gamma=0.1,
    last_epoch=-1,
)

lr = base_lr * gamma ^ floor(step / step_size)

import torch
from torch import nn


model = nn.Linear(10, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer,
    step_size=5,  # 5回step()を呼ぶごとに下げる
    gamma=0.1,    # lrを0.1倍にする
)

for epoch in range(12):
    optimizer.step()   # 実際の訓練ではbackward後に呼ぶ（ここでは形だけ確認するので勾配は使わない）
    scheduler.step()   # schedulerを1回進める。何回目の呼び出しかをschedulerが内部で数えている
    print(epoch, scheduler.get_last_lr())

0 [0.1]
1 [0.1]
2 [0.1]
3 [0.1]
4 [0.010000000000000002]
5 [0.010000000000000002]
6 [0.010000000000000002]
7 [0.010000000000000002]
8 [0.010000000000000002]
9 [0.0010000000000000002]
10 [0.0010000000000000002]
11 [0.0010000000000000002]

net = BanknoteCNN(num_classes=2)  # 第4回で定義したCNN。重み初期化はデフォルトのまま使う

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(
    net.parameters(),
    lr=0.01,
    momentum=0.9,
)
scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer,
    step_size=10,  # 10回のscheduler.step()ごとに下げる（=10 epochごと）
    gamma=0.1,     # 学習率を0.1倍にする
)

for epoch in range(max_epoch):
    net.train()  # BatchNormなどを訓練モードにする

    for images, labels in train_loader:
        outputs = net(images)           # (N, 2) のlogit
        loss = criterion(outputs, labels)

        optimizer.zero_grad(set_to_none=True)  # 前のbatchの勾配を消す
        loss.backward()                        # 勾配を計算する
        optimizer.step()                       # 現在の学習率でパラメータを更新する

    scheduler.step()  # 1epochの訓練ループを抜けたあとで学習率を更新する

lr
|---------
|         |------
|                |----
+---------------------- epoch

torch.optim.lr_scheduler.MultiStepLR(
    optimizer,
    milestones,
    gamma=0.1,
    last_epoch=-1,
)

scheduler = torch.optim.lr_scheduler.MultiStepLR(
    optimizer,
    milestones=[30, 60, 90],  # 30, 60, 90 epochで下げる
    gamma=0.1,
)

import torch
import matplotlib.pyplot as plt  # 可視化のみに使う。未インストールなら pip install matplotlib


weight = torch.randn(1, requires_grad=True)
target = torch.zeros(1)
optimizer = torch.optim.SGD([weight], lr=0.1, momentum=0.9)

milestones = [50, 125, 160]
scheduler = torch.optim.lr_scheduler.MultiStepLR(
    optimizer,
    milestones=milestones,
    gamma=0.1,
)

lr_history = []
epoch_history = []

for epoch in range(200):
    loss = (weight - target).pow(2).sum()  # ダミーのMSE風loss（weightを0へ近づけるだけ）
    loss.backward()
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

    scheduler.step()
    lr_history.append(scheduler.get_last_lr()[0])
    epoch_history.append(epoch)

for epoch in (0, 48, 49, 124, 125, 159, 160, 199):
    print(epoch, lr_history[epoch])

plt.plot(epoch_history, lr_history, label=f"milestones={milestones}")
plt.xlabel("Epoch")
plt.ylabel("Learning rate")
plt.legend()
plt.show()

0 0.1
48 0.1
49 0.010000000000000002
124 0.010000000000000002
125 0.0010000000000000002
159 0.0010000000000000002
160 0.00010000000000000003
199 0.00010000000000000003

torch.optim.lr_scheduler.ExponentialLR(
    optimizer,
    gamma,
    last_epoch=-1,
)

lr_t = base_lr * gamma^t

scheduler = torch.optim.lr_scheduler.ExponentialLR(
    optimizer,
    gamma=0.95,  # stepごとに5%ずつ下げる
)

import torch
import matplotlib.pyplot as plt


weight = torch.randn(1, requires_grad=True)
target = torch.zeros(1)
optimizer = torch.optim.SGD([weight], lr=0.1, momentum=0.9)

gamma = 0.95
scheduler = torch.optim.lr_scheduler.ExponentialLR(
    optimizer,
    gamma=gamma,
)

lr_history = []
epoch_history = []

for epoch in range(200):
    loss = (weight - target).pow(2).sum()
    loss.backward()
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
    scheduler.step()

    lr_history.append(scheduler.get_last_lr()[0])
    epoch_history.append(epoch)

for epoch in (0, 1, 10, 199):
    print(epoch, lr_history[epoch])

plt.plot(epoch_history, lr_history, label=f"gamma={gamma}")
plt.xlabel("Epoch")
plt.ylabel("Learning rate")
plt.legend()
plt.show()

0 0.095
1 0.09025
10 0.05688000922764597
199 3.5052666248828686e-06

torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max,
    eta_min=0.0,
    last_epoch=-1,
)

eta_t = eta_min + 0.5 * (eta_max - eta_min) * (1 + cos(pi * T_cur / T_max))

import torch
import matplotlib.pyplot as plt


weight = torch.randn(1, requires_grad=True)
target = torch.zeros(1)
optimizer = torch.optim.SGD([weight], lr=0.1, momentum=0.9)

t_max = 50
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=t_max,
    eta_min=0.0,
)

lr_history = []
epoch_history = []

for epoch in range(200):
    loss = (weight - target).pow(2).sum()
    loss.backward()
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
    scheduler.step()

    lr_history.append(scheduler.get_last_lr()[0])
    epoch_history.append(epoch)

for epoch in (0, 25, 49, 50, 75, 99, 100, 199):
    print(epoch, lr_history[epoch])

plt.plot(epoch_history, lr_history, label=f"T_max={t_max}")
plt.xlabel("Epoch")
plt.ylabel("Learning rate")
plt.legend()
plt.show()

0 0.09990133642141358
25 0.046860474023534326
49 0.0
50 9.866357858642206e-05
75 0.05313952597646564
99 0.10000000000000002
100 0.09990133642141359
199 0.10000000000000005

torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer,
    T_0,
    T_mult=1,
    eta_min=0.0,
    last_epoch=-1,
)

torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode="min",
    factor=0.1,
    patience=10,
    threshold=1e-4,
    threshold_mode="rel",
    cooldown=0,
    min_lr=0,
    eps=1e-8,
)

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode="min",
    factor=0.5,
    patience=3,
)

for epoch in range(num_epochs):
    train_one_epoch(...)
    valid_loss = evaluate(...)
    scheduler.step(valid_loss)  # 検証が終わってからmetricを渡す

import torch


weight = torch.randn(1, requires_grad=True)
optimizer = torch.optim.SGD([weight], lr=0.1, momentum=0.9)

valid_loss = 0.5
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode="min",
    factor=0.1,
    patience=10,
    cooldown=10,
    min_lr=1e-4,
)

for epoch in range(50):
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

    if epoch == 5:
        valid_loss = 0.4  # ここだけ改善したとみなす

    scheduler.step(valid_loss)
    if epoch in (0, 15, 16, 26, 27, 37, 49):
        print(epoch, optimizer.param_groups[0]["lr"])

0 0.1
15 0.1
16 0.010000000000000002
26 0.010000000000000002
27 0.010000000000000002
37 0.0010000000000000002
49 0.0010000000000000002

torch.optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda,
    last_epoch=-1,
)

def factor(epoch):
    return 0.95 ** epoch  # base_lrに掛ける倍率


scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda=factor,
)

weights_1 = torch.randn(6, 3, 5, 5, requires_grad=True)  # 1つ目のparam group（畳み込みカーネル風の形）
weights_2 = torch.ones(5, 5, requires_grad=True)         # 2つ目のparam group

optimizer = torch.optim.SGD(
    [
        {"params": [weights_1]},  # group 0。lambda1で倍率を決める
        {"params": [weights_2]},  # group 1。lambda2で倍率を決める
    ],
    lr=0.1,
)

def lambda1(epoch):
    return 0.1 ** (epoch // 20)  # 20 epochごとに0.1倍（0.1 -> 0.01 -> 0.001）


def lambda2(epoch):
    return 0.95 ** epoch  # 毎epoch0.95倍


scheduler = torch.optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda=[lambda1, lambda2],  # リストの順番がparam groupの順番に対応する
)

for epoch in range(60):
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
    scheduler.step()
    if epoch in (0, 1, 19, 20, 39, 40, 59):
        print(epoch, scheduler.get_last_lr())

0 [0.1, 0.095]
1 [0.1, 0.09025]
19 [0.010000000000000002, 0.03584859224085419]
20 [0.010000000000000002, 0.03405616262881148]
39 [0.0010000000000000002, 0.012851215656510312]
40 [0.0010000000000000002, 0.012208654873684797]
59 [0.00010000000000000003, 0.004606979898695194]

import torch


weight = torch.randn(1, requires_grad=True)
optimizer = torch.optim.SGD([weight], lr=0.1, momentum=0.9)

warmup = torch.optim.lr_scheduler.LinearLR(
    optimizer,
    start_factor=0.1,  # 最初はbase_lrの10%から始める
    total_iters=5,     # 5回かけてbase_lrへ近づける
)
cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=95,  # warmup後に残る95epoch分で、base_lrから0まで下げる
)
scheduler = torch.optim.lr_scheduler.SequentialLR(
    optimizer,
    schedulers=[warmup, cosine],
    milestones=[5],  # 5回stepしたらcosineへ切り替える
)

lr_history = []
for epoch in range(100):
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)
    scheduler.step()
    lr_history.append(scheduler.get_last_lr()[0])

for epoch in (0, 1, 2, 3, 4, 5, 6, 50, 99):
    print(epoch, lr_history[epoch])

0 0.028000000000000004
1 0.046000000000000006
2 0.064
3 0.082
4 0.1
5 0.09997266286704631
6 0.09989068136093875
50 0.05247918773366114
99 0.0

torch.optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr,
    epochs=None,
    steps_per_epoch=None,
    total_steps=None,
)

torch.optim.lr_scheduler.CyclicLR(
    optimizer,
    base_lr,
    max_lr,
    step_size_up=2000,
    step_size_down=None,
)

torch.optim.lr_scheduler.SequentialLR(
    optimizer,
    schedulers,
    milestones,
)

torch.optim.lr_scheduler.ChainedScheduler(schedulers)

checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "epoch": epoch,
    "global_step": global_step,
}

torch.save(checkpoint, "checkpoint.pt")

checkpoint = torch.load("checkpoint.pt", map_location=device, weights_only=True)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
scheduler.load_state_dict(checkpoint["scheduler"])
start_epoch = checkpoint["epoch"] + 1
global_step = checkpoint["global_step"]

PyTorchの訓練スクリプト
  -> SummaryWriter
  -> event fileをlog_dirへ書き込む
  -> tensorboard --logdir runs
  -> ブラウザでダッシュボードを見る

python -m pip install tensorboard

tensorboard --version
tensorboard --logdir runs --port 6006

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(
    log_dir=None,
    comment="",
    purge_step=None,
    max_queue=10,
    flush_secs=120,
    filename_suffix="",
)

from torch.utils.tensorboard import SummaryWriter


with SummaryWriter(log_dir="runs/exp01") as writer:
    writer.add_scalar("Loss/train", 0.5, 1)

writer.flush()
writer.close()

import numpy as np
from torch.utils.tensorboard import SummaryWriter


writer = SummaryWriter(comment="test_tensorboard")

for step in range(100):
    writer.add_scalar("y=2x", step * 2, step)
    writer.add_scalar("y=pow(2, x)", 2 ** step, step)
    writer.add_scalars(
        "data/scalar_group",
        {
            "xsinx": step * np.sin(step),
            "xcosx": step * np.cos(step),
            "arctanx": np.arctan(step),
        },
        step,
    )

writer.close()

writer.add_scalar(
    tag,
    scalar_value,
    global_step=None,
    walltime=None,
    new_style=False,
    double_precision=False,
)

writer.add_scalar("Loss/train", train_loss, epoch)
writer.add_scalar("Accuracy/valid", valid_accuracy, epoch)
writer.add_scalar("LearningRate/group0", optimizer.param_groups[0]["lr"], epoch)

global_step = 0

for epoch in range(max_epoch):
    net.train()

    for images, labels in train_loader:
        global_step += 1

        outputs = net(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

        predicted = outputs.argmax(dim=1)
        batch_accuracy = (predicted == labels).float().mean().item()

        writer.add_scalar("Loss/train_step", loss.item(), global_step)
        writer.add_scalar("Accuracy/train_step", batch_accuracy, global_step)

writer.add_scalars(
    main_tag,
    tag_scalar_dict,
    global_step=None,
    walltime=None,
)

writer.add_scalars(
    "Loss/epoch",
    {
        "train": train_loss,
        "valid": valid_loss,
    },
    epoch,
)

for step in range(100):
    writer.add_scalars(
        "data/scalar_group",
        {
            "xsinx": step * np.sin(step),
            "xcosx": step * np.cos(step),
        },
        step,
    )

writer.add_histogram(
    tag,
    values,
    global_step=None,
    bins="tensorflow",
    walltime=None,
    max_bins=None,
)

for name, parameter in model.named_parameters():
    writer.add_histogram(f"Parameters/{name}", parameter.detach(), epoch)
    if parameter.grad is not None:
        writer.add_histogram(f"Gradients/{name}", parameter.grad.detach(), epoch)

import numpy as np
from torch.utils.tensorboard import SummaryWriter


writer = SummaryWriter(comment="histogram_demo")

for step in range(2):
    np.random.seed(step)
    uniform_values = np.arange(100)
    normal_values = np.random.normal(size=1000)

    writer.add_histogram("distribution/uniform", uniform_values, step)
    writer.add_histogram("distribution/normal", normal_values, step)

writer.close()

for epoch in range(max_epoch):
    train_one_epoch(...)

    for name, parameter in net.named_parameters():
        writer.add_histogram(f"Weights/{name}", parameter.detach(), epoch)
        if parameter.grad is not None:
            writer.add_histogram(f"Gradients/{name}", parameter.grad.detach(), epoch)

writer.add_image(
    tag,
    img_tensor,
    global_step=None,
    walltime=None,
    dataformats="CHW",
)

import torch
from torch.utils.tensorboard import SummaryWriter


writer = SummaryWriter(comment="image_demo")

chw_image = torch.rand(3, 512, 512)
writer.add_image("fake_img/chw", chw_image, 1)

hw_image = torch.rand(512, 512)
writer.add_image("fake_img/hw", hw_image, 2, dataformats="HW")

hwc_image = torch.rand(512, 512, 3)
writer.add_image("fake_img/hwc", hwc_image, 3, dataformats="HWC")

writer.close()

writer.add_images(tag, img_tensor, global_step=None, walltime=None, dataformats="NCHW")

from torchvision.utils import make_grid

grid = make_grid(
    images[:16],
    nrow=4,
    padding=2,
    normalize=True,
    value_range=None,
    scale_each=False,
    pad_value=0.0,
)
writer.add_image("Batch/train", grid, global_step)

from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from torchvision.utils import make_grid


writer = SummaryWriter(comment="make_grid_demo")

transform = transforms.Compose([
    transforms.Resize((32, 64)),
    transforms.ToTensor(),
])
train_data = BanknoteDataset(data_dir=train_dir, transform=transform)
train_loader = DataLoader(train_data, batch_size=16, shuffle=True)

images, labels = next(iter(train_loader))
image_grid = make_grid(
    images,
    nrow=4,
    normalize=True,
    scale_each=True,
)
writer.add_image("input_images", image_grid, 0)
writer.close()

writer.add_graph(
    model,
    input_to_model=None,
    verbose=False,
    use_strict_trace=True,
)

example = torch.randn(1, 3, 224, 224, device=device)
writer.add_graph(model, example)

from torch import nn
from torch.utils.tensorboard import SummaryWriter


writer = SummaryWriter(comment="graph_demo")

fake_image = torch.randn(1, 3, 32, 32)
model = nn.Sequential(
    nn.Conv2d(3, 16, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Flatten(),
    nn.Linear(16 * 16 * 16, 2),  # 32x32がMaxPool2d(2)で16x16に、チャンネル数は16
)

writer.add_graph(model, fake_image)
writer.close()

model = BanknoteCNN(num_classes=2)  # 第4回で定義したCNN

print(model)

total = 0
for name, parameter in model.named_parameters():
    count = parameter.numel()
    total += count
    print(name, parameter.shape, count)

print("total parameters:", total)

features.0.weight torch.Size([32, 3, 3, 3]) 864
features.1.weight torch.Size([32]) 32
features.1.bias torch.Size([32]) 32
features.4.weight torch.Size([64, 32, 3, 3]) 18432
features.5.weight torch.Size([64]) 64
features.5.bias torch.Size([64]) 64
features.8.weight torch.Size([128, 64, 3, 3]) 73728
features.9.weight torch.Size([128]) 128
features.9.bias torch.Size([128]) 128
classifier.weight torch.Size([2, 128]) 256
classifier.bias torch.Size([2]) 2
total parameters: 93730

from torchvision.models.feature_extraction import create_feature_extractor


return_nodes = {
    "features.2": "pool1",
    "features.5": "pool2",
}
extractor = create_feature_extractor(model, return_nodes=return_nodes)
features = extractor(x)

weights = model.features[0].weight.detach().cpu()

import torch.nn as nn
import torchvision.models as models
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import make_grid


writer = SummaryWriter(comment="kernel_demo")
alexnet = models.alexnet(weights=models.AlexNet_Weights.DEFAULT)

for layer_index, module in enumerate(alexnet.features):
    if isinstance(module, nn.Conv2d):
        kernels = module.weight.detach().cpu()
        out_channels, in_channels, kernel_h, kernel_w = kernels.shape

        if in_channels == 3:
            kernel_grid = make_grid(
                kernels,
                nrow=8,
                normalize=True,
                scale_each=True,
            )
            writer.add_image(f"conv{layer_index}/kernels", kernel_grid, 0)
        break

writer.close()

from PIL import Image
from torchvision import transforms


image = Image.open("your_image.jpg").convert("RGB")  # 任意のRGB画像に置き換える
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.49139968, 0.48215827, 0.44653124],
        std=[0.24703233, 0.24348505, 0.26158768],
    ),
])
image_tensor = transform(image).unsqueeze(0)

first_conv = alexnet.features[0]
feature_map = first_conv(image_tensor).detach().cpu()
feature_map = feature_map.transpose(0, 1)  # (1, C, H, W) -> (C, 1, H, W)
feature_grid = make_grid(feature_map, normalize=True, scale_each=True, nrow=8)

writer = SummaryWriter(comment="feature_map_demo")
writer.add_image("feature_map/conv1", feature_grid, 0)
writer.close()

Module呼び出し
  -> forward_pre_hook
  -> forward
  -> forward_hook
  -> loss
  -> backward
  -> Tensor gradient hook / full backward hook

handle = module.register_forward_hook(hook)
...
handle.remove()

activations[name] = output.detach().cpu()

handle = tensor.register_hook(hook)

hook(grad) -> Tensor または None

import torch


x = torch.tensor(2.0, requires_grad=True)
y = x ** 2  # y = x^2 なので dy/dx = 2x

def print_grad(grad):
    print("grad:", grad)
    return None


handle = x.register_hook(print_grad)
y.backward()
handle.remove()

grad: tensor(4.)

import torch


w = torch.tensor([1.0], requires_grad=True)
x = torch.tensor([2.0], requires_grad=True)

a = w + x  # a = 1.0 + 2.0 = 3.0
b = w + 1  # b = 1.0 + 1.0 = 2.0
y = a * b  # y = 3.0 * 2.0 = 6.0

a_grad = []

def save_grad(grad):
    a_grad.append(grad.detach().clone())  # 勾配を書き換えず、コピーを保存するだけ


handle = a.register_hook(save_grad)
y.backward()

print("w.grad:", w.grad)
print("x.grad:", x.grad)
print("a_grad:", a_grad[0])

handle.remove()

w.grad: tensor([5.])
x.grad: tensor([2.])
a_grad: tensor([2.])

w = torch.tensor([1.0], requires_grad=True)
x = torch.tensor([2.0], requires_grad=True)

y = (w + x) * (w + 1)  # 前の例と同じ式（Hookがなければw.gradは5になるはず）

def scale_grad(grad):
    return grad * 3  # 返したTensorで、この場所を流れる勾配を置き換える


handle = w.register_hook(scale_grad)
y.backward()

print("w.grad:", w.grad)
handle.remove()

w.grad: tensor([15.])

module.register_forward_hook(
    hook,
    *,
    prepend=False,
    with_kwargs=False,
    always_call=False,
)

hook(module, args, output) -> None または新しい output

activations = {}

def save_activation(name):
    def hook(module, args, output):
        activations[name] = output.detach().cpu()
    return hook


handle = model.features[-1].register_forward_hook(
    save_activation("last_conv")
)

_ = model(images)
handle.remove()

import math
import torch
import torch.nn as nn
import torchvision.models as models
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import make_grid


writer = SummaryWriter(comment="hook_feature_map_demo")
model = models.alexnet(weights=models.AlexNet_Weights.DEFAULT).eval()

# 実際には前節で説明したPIL+transformsの前処理を通した本物の画像を使う。
# ここではHook登録の仕組みそのものを確認したいので、ダミーの1枚のRGB画像で代用する。
image_tensor = torch.randn(1, 3, 224, 224)

feature_maps = {}
handles = []

def save_feature_map(name):
    def hook(module, args, output):
        feature_maps[name] = output.detach().cpu()
    return hook


for name, module in model.named_modules():
    if isinstance(module, nn.Conv2d):
        handle = module.register_forward_hook(save_feature_map(name))
        handles.append(handle)

_ = model(image_tensor)  # image_tensorのshapeは (1, 3, 224, 224)

for name, feature_map in feature_maps.items():
    channels = feature_map.transpose(0, 1)  # (1, C, H, W) -> (C, 1, H, W)
    nrow = max(1, int(math.sqrt(channels.shape[0])))
    grid = make_grid(channels, normalize=True, scale_each=True, nrow=nrow)
    writer.add_image(f"feature_map/{name}", grid, 0)

for handle in handles:
    handle.remove()

writer.close()

module.register_forward_pre_hook(
    hook,
    *,
    prepend=False,
    with_kwargs=False,
)

hook(module, args) -> None または新しい args

import torch
from torch import nn


class TinyConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 2, kernel_size=3)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.pool1(x)
        return x


def forward_hook(module, args, output):
    print("forward_hook output shape:", output.shape)


def forward_pre_hook(module, args):
    print("forward_pre_hook input shape:", args[0].shape)


def backward_hook(module, grad_input, grad_output):
    print("backward_hook grad_output shape:", grad_output[0].shape)


net = TinyConvNet()
net.conv1.weight.data[0].fill_(1.0)  # 出力チャンネル0のカーネルをすべて1.0にする
net.conv1.weight.data[1].fill_(2.0)  # 出力チャンネル1のカーネルをすべて2.0にする
net.conv1.bias.data.zero_()

handles = [
    net.conv1.register_forward_pre_hook(forward_pre_hook),
    net.conv1.register_forward_hook(forward_hook),
    net.conv1.register_full_backward_hook(backward_hook),
]

fake_image = torch.ones(1, 1, 4, 4)  # 全要素が1.0の4x4画像1枚
output = net(fake_image)
target = torch.randn_like(output)
loss = nn.L1Loss()(output, target)
loss.backward()

for handle in handles:
    handle.remove()

forward_pre_hook input shape: torch.Size([1, 1, 4, 4])
forward_hook output shape: torch.Size([1, 2, 2, 2])
backward_hook grad_output shape: torch.Size([1, 2, 2, 2])

module.register_full_backward_pre_hook(hook)
module.register_full_backward_hook(hook)

hook(module, grad_input, grad_output) -> None または新しい grad_input

最後の畳み込み特徴 A_k(x, y)
  -> Global Average Poolingで F_k にする
  -> Linear分類器の重み w_k^c を使う
  -> クラスcのスコア y^c = sum_k w_k^c F_k

M_c(x, y) = sum_k w_k^c * A_k(x, y)
Heatmap = ReLU(M_c)

alpha_k^c = (1 / Z) * sum_i sum_j d y^c / d A_ij^k

L_Grad-CAM^c = ReLU(sum_k alpha_k^c * A^k)

import torch
import torch.nn.functional as F


class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.activations = None
        self.gradients = None
        self.forward_handle = target_layer.register_forward_hook(
            self._forward_hook
        )

    def _forward_hook(self, module, args, output):
        self.activations = output                  # 対象層の特徴マップを保存する
        output.register_hook(self._gradient_hook)  # この特徴マップへの勾配を保存する

    def _gradient_hook(self, grad):
        self.gradients = grad

    def __call__(self, inputs, class_index=None):
        self.model.zero_grad(set_to_none=True)
        logits = self.model(inputs)  # shapeは通常 (N, num_classes)

        if class_index is None:
            class_index = logits.argmax(dim=1)  # 指定がなければ予測クラスを見る

        score = logits.gather(1, class_index[:, None]).sum()
        score.backward()

        weights = self.gradients.mean(dim=(-2, -1), keepdim=True)
        cam = (weights * self.activations).sum(dim=1, keepdim=True)
        cam = F.relu(cam)
        cam = F.interpolate(
            cam,
            size=inputs.shape[-2:],
            mode="bilinear",
            align_corners=False,
        )

        minimum = cam.amin(dim=(-2, -1), keepdim=True)
        maximum = cam.amax(dim=(-2, -1), keepdim=True)
        cam = (cam - minimum) / (maximum - minimum).clamp_min(1e-8)
        return logits, cam.detach()

    def close(self):
        self.forward_handle.remove()

class TinyCAMNet(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 4, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(4, 8, kernel_size=3, padding=1)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(8, num_classes)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))  # conv2の出力をGrad-CAMの対象にする
        x = self.pool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x


torch.manual_seed(0)
model = TinyCAMNet(num_classes=2)
inputs = torch.randn(2, 3, 16, 16)  # batch size 2、3チャンネル、16x16のダミー画像

cam_extractor = GradCAM(model, model.conv2)
logits, heatmap = cam_extractor(inputs)

print("logits.shape:", logits.shape)
print("heatmap.shape:", heatmap.shape)
print("heatmap.min():", heatmap.min().item(), "heatmap.max():", heatmap.max().item())

cam_extractor.close()

logits.shape: torch.Size([2, 2])
heatmap.shape: torch.Size([2, 1, 16, 16])
heatmap.min(): 0.0 heatmap.max(): 1.0

訓練ループ
  ├── add_scalar: loss / metric / lr
  ├── add_histogram: parameter / gradient
  ├── add_image: input / prediction / heatmap
  └── add_graph: model trace
        |
        v
SummaryWriter -> event file -> TensorBoard -> ブラウザ表示

入力 args / kwargs
    |
forward_pre_hook
    |
module.forward
    |
forward_hook
    |
loss
    |
backward
    |
Tensor gradient hook / full backward hook

対象クラスのlogit
  -> 対象層特徴マップへの勾配を計算
  -> 勾配を空間平均してチャンネル重み alpha を作る
  -> 特徴マップを重み付き和にする
  -> ReLU
  -> 入力画像サイズへリサイズ
  -> 正規化して元画像へ重ねる

from pathlib import Path
from torch.utils.tensorboard import SummaryWriter


log_dir = Path("runs") / "experiment_001"
writer = SummaryWriter(log_dir=str(log_dir))

global_step = 0

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0

    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)
        logits = model(images)
        loss = criterion(logits, targets)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        writer.add_scalar("Loss/train_step", loss.item(), global_step)
        writer.add_scalar(
            "LearningRate/group0",
            optimizer.param_groups[0]["lr"],
            global_step,
        )
        global_step += 1

        # OneCycleLRのようにbatch単位で進めるschedulerならここでstepする
        # scheduler.step()

    train_loss = running_loss / len(train_loader)
    valid_loss, valid_accuracy = evaluate(model, valid_loader)

    writer.add_scalars(
        "Loss/epoch",
        {
            "train": train_loss,
            "valid": valid_loss,
        },
        epoch,
    )
    writer.add_scalar("Accuracy/valid", valid_accuracy, epoch)

    if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
        scheduler.step(valid_loss)
    else:
        scheduler.step()

    if epoch % 5 == 0:
        for name, parameter in model.named_parameters():
            writer.add_histogram(f"Parameters/{name}", parameter.detach(), epoch)
            if parameter.grad is not None:
                writer.add_histogram(
                    f"Gradients/{name}",
                    parameter.grad.detach(),
                    epoch,
                )

writer.close()