PyTorch入門 第7回：正則化、Weight Decay、Dropout、Normalizationで過学習を抑える

第6回：訓練を制御・記録・観察する
  -> lossやaccuracyの動きが分かる
  -> 学習率や勾配の状態が見える

第7回：過学習や不安定さを抑える
  -> Weight Decayでパラメータを制御する
  -> Dropoutで活性値へランダム性を入れる
  -> Normalizationで特徴の尺度を整える

汎化と訓練安定性を助ける方法
├── パラメータへ直接働く方法
│   ├── L1正則化
│   └── L2正則化 / Weight Decay
├── 活性値へランダム性を入れる方法
│   └── Dropout
└── 特徴の平均・分散・尺度を整える方法
    ├── BatchNorm
    ├── LayerNorm
    ├── InstanceNorm
    └── GroupNorm

Expected Error = Bias^2 + Variance + Irreducible Noise

R_L1(w) = sum_i |w_i|
Objective = DataLoss + lambda * sum_i |w_i|

import torch


# requires_grad=Trueにして、このTensor自体に対するペナルティ（L1項）を計算できるようにする
weight = torch.tensor(
    [[1.0, -2.0, 0.5], [0.0, 0.25, -0.75]],
    requires_grad=True,
)
bias = torch.tensor([0.1, -0.2], requires_grad=True)

# 全要素の絶対値を足し合わせたものがL1正則化項
l1_penalty = weight.abs().sum() + bias.abs().sum()
print(round(l1_penalty.item(), 4))
# 4.8

objective = data_loss + l1_lambda * l1_penalty

R_L2(w) = (1/2) * sum_i w_i^2
Objective = DataLoss + (lambda / 2) * ||w||_2^2

∂R_L2 / ∂w = w
∂Objective / ∂w = ∂DataLoss / ∂w + lambda * w

w_(t+1)
= w_t - eta * (grad DataLoss(w_t) + lambda * w_t)
= (1 - eta * lambda) * w_t - eta * grad DataLoss(w_t)

import torch


w = torch.tensor([2.0], requires_grad=True)
# weight_decay=0.1は、勾配へ lambda * w を足す形でWeight Decayを実現する引数
optimizer = torch.optim.SGD([w], lr=0.1, weight_decay=0.1)

loss = 0.5 * (w - 5).pow(2).sum()  # wを5に近づけたいだけの単純な二乗誤差
loss.backward()

print("grad before weight_decay:", round(w.grad.item(), 4))
# grad before weight_decay: -3.0
optimizer.step()  # 内部でWeight Decay分（lambda * w = 0.2）が勾配に足されてから更新される
print("weight after step:", round(w.item(), 4))
# weight after step: 2.28

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


torch.manual_seed(1)

def gen_data(num_data=10, x_range=(-1, 1)):
    w = 1.5
    train_x = torch.linspace(*x_range, num_data).unsqueeze(1)
    train_y = w * train_x + torch.normal(0, 0.5, size=train_x.size())
    test_x = torch.linspace(*x_range, num_data).unsqueeze(1)
    test_y = w * test_x + torch.normal(0, 0.3, size=test_x.size())
    return train_x, train_y, test_x, test_y


class TinyMLP(nn.Module):
    def __init__(self, hidden=64):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(1, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, 1),
        )

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


def weight_l2_norm(model):
    # 全パラメータのうちweightという名前が付くものだけを対象に、L2ノルム（大きさ）を1つの数値にまとめる
    total = torch.tensor(0.0)
    for name, parameter in model.named_parameters():
        if "weight" in name:
            total = total + parameter.detach().pow(2).sum()
    return total.sqrt().item()


train_x, train_y, test_x, test_y = gen_data()  # 訓練10点・検証10点の1次元データ

net_normal = TinyMLP()
net_wdecay = copy.deepcopy(net_normal)  # 初期値をそろえてweight_decayだけを比較する

optimizer_normal = torch.optim.SGD(
    net_normal.parameters(),
    lr=0.01,
    momentum=0.9,
)
optimizer_wdecay = torch.optim.SGD(
    net_wdecay.parameters(),
    lr=0.01,
    momentum=0.9,
    weight_decay=1e-2,  # net_normal用のoptimizerとの違いはこの1行だけ
)
loss_fn = nn.MSELoss()
writer = SummaryWriter(log_dir="runs/weight_decay_demo")

for epoch in range(400):
    pred_normal = net_normal(train_x)
    pred_wdecay = net_wdecay(train_x)
    loss_normal = loss_fn(pred_normal, train_y)
    loss_wdecay = loss_fn(pred_wdecay, train_y)

    optimizer_normal.zero_grad(set_to_none=True)
    optimizer_wdecay.zero_grad(set_to_none=True)
    loss_normal.backward()
    loss_wdecay.backward()
    optimizer_normal.step()
    optimizer_wdecay.step()

    if epoch + 1 in (1, 200, 400):
        with torch.no_grad():
            test_loss_normal = loss_fn(net_normal(test_x), test_y)
            test_loss_wdecay = loss_fn(net_wdecay(test_x), test_y)
        print(
            epoch + 1,
            round(loss_normal.item(), 4),
            round(loss_wdecay.item(), 4),
            round(test_loss_normal.item(), 4),
            round(test_loss_wdecay.item(), 4),
        )
        # epoch, train loss（通常）, train loss（weight_decay）, test loss（通常）, test loss（weight_decay）の順
        # 1 0.6151 0.6151 0.8963 0.8963
        # 200 0.0538 0.0563 0.0829 0.0835
        # 400 0.0448 0.0525 0.0894 0.0862

    if (epoch + 1) % 200 == 0:
        # weightと勾配の分布はヒストグラム画像としてTensorBoardに記録される（テキストでは表示できない）
        for name, parameter in net_normal.named_parameters():
            writer.add_histogram(f"normal/{name}", parameter.detach(), epoch + 1)
            if parameter.grad is not None:
                writer.add_histogram(f"normal_grad/{name}", parameter.grad, epoch + 1)
        for name, parameter in net_wdecay.named_parameters():
            writer.add_histogram(f"weight_decay/{name}", parameter.detach(), epoch + 1)
            if parameter.grad is not None:
                writer.add_histogram(f"weight_decay_grad/{name}", parameter.grad, epoch + 1)

print("weight_norm:", round(weight_l2_norm(net_normal), 4), round(weight_l2_norm(net_wdecay), 4))
# weight_norm: 8.0911 5.6566
writer.close()

g_t = grad DataLoss(w_t) + lambda * w_t
optimizerは g_t を使って更新する

w_t <- (1 - eta * lambda) * w_t
その後、Adamの自適応勾配更新を行う

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-3,
    weight_decay=1e-2,
)

import torch
from torch import nn


model = nn.Sequential(
    nn.Conv2d(3, 8, kernel_size=3, padding=1, bias=False),
    nn.BatchNorm2d(8),
    nn.Flatten(),
    nn.Linear(8 * 4 * 4, 2),
)

decay_names = []
no_decay_names = []
decay = []
no_decay = []

for name, parameter in model.named_parameters():
    if not parameter.requires_grad:
        continue

    if parameter.ndim == 1 or name.endswith("bias"):
        # BatchNormのweight/bias、Linearのbiasなど1次元パラメータはWeight Decayから外す
        no_decay_names.append(name)
        no_decay.append(parameter)
    else:
        decay_names.append(name)
        decay.append(parameter)

optimizer = torch.optim.AdamW(
    [
        {"params": decay, "weight_decay": 0.01},
        {"params": no_decay, "weight_decay": 0.0},
    ],
    lr=1e-3,
)

print("decay:", decay_names)
# decay: ['0.weight', '3.weight']
print("no_decay:", no_decay_names)
# no_decay: ['1.weight', '1.bias', '3.bias']

m_i ~ Bernoulli(1 - p)
y_i = m_i * x_i / (1 - p)

train: y = mask * x / (1 - p)
eval : y = x

nn.Dropout(p=0.5, inplace=False)

import torch
from torch import nn


torch.manual_seed(0)  # マスクのサンプリングを再現できるようにする
x = torch.ones(10)
dropout = nn.Dropout(p=0.5)

dropout.train()  # 訓練モード：要素を確率的に0にし、残った値を1/(1-p)倍する
y_train = dropout(x)

dropout.eval()  # 評価モード：恒等写像になり、何も変えない
y_eval = dropout(x)

print("train:", y_train)
# train: tensor([0., 0., 2., 0., 0., 0., 2., 2., 0., 2.])
print("eval:", y_eval)
# eval: tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])
print("zeros:", int((y_train == 0).sum()))
# zeros: 6

import torch
from torch import nn


torch.manual_seed(1)

class DropoutLinear(nn.Module):
    def __init__(self, input_num, d_prob=0.5):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Dropout(d_prob),
            nn.Linear(input_num, 1, bias=False),
            nn.ReLU(),
        )

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


input_num = 10000
x = torch.ones(input_num, dtype=torch.float32)

net = DropoutLinear(input_num, d_prob=0.5)
net.layers[1].weight.detach().fill_(1.0)  # 重みをすべて1にして、和がそのまま出力になるようにする

net.train()
print("training:", net(x))
# training: tensor([9890.], grad_fn=<ReluBackward0>)

net.eval()
print("eval:", net(x))
# eval: tensor([10000.], grad_fn=<ReluBackward0>)

import torch
from torch import nn


layer = nn.Dropout(p=0.5)
input_tensor = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

layer.eval()  # 評価モードなのでDropoutは恒等写像になる
output = layer(input_tensor).sum()
output.backward()

print("grad:", input_tensor.grad)
# grad: tensor([1., 1., 1.])

import torch
from torch import nn


torch.manual_seed(0)
x = torch.ones(1, 3, 2, 2)  # (N, C, H, W) = (1枚, 3チャンネル, 2x2)
dropout2d = nn.Dropout2d(p=0.5)
dropout2d.train()

y = dropout2d(x)
print(y)
# tensor([[[[0., 0.],
#           [0., 0.]],
#
#          [[0., 0.],
#           [0., 0.]],
#
#          [[2., 2.],
#           [2., 2.]]]])
print("channel sums:", y.sum(dim=(0, 2, 3)))
# channel sums: tensor([0., 0., 8.])

import copy
import torch
from torch import nn


torch.manual_seed(1)

def gen_data(num_data=10, x_range=(-1, 1)):
    w = 1.5
    train_x = torch.linspace(*x_range, num_data).unsqueeze(1)
    train_y = w * train_x + torch.normal(0, 0.5, size=train_x.size())
    test_x = torch.linspace(*x_range, num_data).unsqueeze(1)
    test_y = w * test_x + torch.normal(0, 0.3, size=test_x.size())
    return train_x, train_y, test_x, test_y


class DropoutMLP(nn.Module):
    def __init__(self, hidden=64, d_prob=0.5):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(1, hidden),
            nn.ReLU(),
            nn.Dropout(d_prob),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Dropout(d_prob),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Dropout(d_prob),
            nn.Linear(hidden, 1),
        )

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


train_x, train_y, test_x, test_y = gen_data()
net_p0 = DropoutMLP(d_prob=0.0)
net_p05 = copy.deepcopy(net_p0)  # 同じ初期値から始め、d_probだけを変える
for module in net_p05.modules():
    if isinstance(module, nn.Dropout):
        module.p = 0.5

optimizer_p0 = torch.optim.SGD(net_p0.parameters(), lr=0.01, momentum=0.9)
optimizer_p05 = torch.optim.SGD(net_p05.parameters(), lr=0.01, momentum=0.9)
loss_fn = nn.MSELoss()

for epoch in range(600):
    net_p0.train()
    net_p05.train()
    pred_p0 = net_p0(train_x)
    pred_p05 = net_p05(train_x)
    loss_p0 = loss_fn(pred_p0, train_y)
    loss_p05 = loss_fn(pred_p05, train_y)

    optimizer_p0.zero_grad(set_to_none=True)
    optimizer_p05.zero_grad(set_to_none=True)
    loss_p0.backward()
    loss_p05.backward()
    optimizer_p0.step()
    optimizer_p05.step()

    if epoch + 1 in (200, 400, 600):
        net_p0.eval()  # 検証lossを見るときはDropoutを恒等写像に戻す
        net_p05.eval()
        with torch.no_grad():
            test_p0 = loss_fn(net_p0(test_x), test_y)
            test_p05 = loss_fn(net_p05(test_x), test_y)
        print(
            epoch + 1,
            round(loss_p0.item(), 4),
            round(loss_p05.item(), 4),
            round(test_p0.item(), 4),
            round(test_p05.item(), 4),
        )
        # epoch, train loss（Dropoutなし）, train loss（Dropoutあり）, test loss（Dropoutなし）, test loss（Dropoutあり）の順
        # 200 0.0538 0.2197 0.0829 0.2512
        # 400 0.0448 0.0789 0.0894 0.2802
        # 600 0.0337 0.1308 0.0993 0.305

mu_B = (1 / m) * sum_i x_i
sigma_B^2 = (1 / m) * sum_i (x_i - mu_B)^2
x_hat_i = (x_i - mu_B) / sqrt(sigma_B^2 + eps)
y_i = gamma * x_hat_i + beta

import torch
from torch import nn


torch.manual_seed(1)

class DeepMLP(nn.Module):
    def __init__(self, width=64, layers=12, use_bn=True):
        super().__init__()
        self.linears = nn.ModuleList(
            [nn.Linear(width, width, bias=False) for _ in range(layers)]
        )
        self.bns = nn.ModuleList(
            [nn.BatchNorm1d(width) for _ in range(layers)]
        ) if use_bn else None
        self.use_bn = use_bn

        for linear in self.linears:
            nn.init.normal_(linear.weight, mean=0.0, std=1.0)  # あえて大きめの標準偏差で初期化する

    def forward(self, x):
        stds = []
        for index, linear in enumerate(self.linears):
            x = linear(x)
            if self.use_bn:
                x = self.bns[index](x)  # BatchNormで各層の分布を毎回整え直す
            x = torch.relu(x)
            stds.append(round(x.std().item(), 4))  # その層の出力スケールを記録する
        return x, stds


inputs = torch.randn(16, 64)

for use_bn in (False, True):
    model = DeepMLP(use_bn=use_bn)
    _, stds = model(inputs)
    print("use_bn=", use_bn, "first5=", stds[:5], "last=", stds[-1])
    # use_bn= False first5= [5.0433, 27.9425, 144.6259, 789.2551, 4464.2241] last= 628074304.0
    # use_bn= True first5= [0.5718, 0.5771, 0.5668, 0.5868, 0.5663] last= 0.5722

import torch
from torch import nn


bn = nn.BatchNorm1d(3, affine=True, track_running_stats=True)
with torch.no_grad():
    bn.weight.fill_(1.0)  # gamma=1、beta=0にして、標準化の結果だけを見やすくする
    bn.bias.zero_()

x = torch.tensor(
    [[1.0, 2.0, 3.0],
     [3.0, 4.0, 5.0]]
)

y = bn(x)  # 訓練モード（既定）なので、このmini-batch自身の平均・分散で正規化する

print("output:", torch.round(y, decimals=4))
# output: tensor([[-1., -1., -1.],
#         [ 1.,  1.,  1.]], grad_fn=<RoundBackward1>)
print("running_mean:", torch.round(bn.running_mean, decimals=4))
# running_mean: tensor([0.2000, 0.3000, 0.4000])
print("running_var:", torch.round(bn.running_var, decimals=4))
# running_var: tensor([1.1000, 1.1000, 1.1000])

running_new = (1 - momentum) * running_old + momentum * batch_stat

nn.BatchNorm1d(num_features, eps=1e-5, momentum=0.1,
               affine=True, track_running_stats=True,
               device=None, dtype=None)

nn.BatchNorm2d(num_features, eps=1e-5, momentum=0.1,
               affine=True, track_running_stats=True,
               device=None, dtype=None)

nn.BatchNorm3d(num_features, eps=1e-5, momentum=0.1,
               affine=True, track_running_stats=True,
               device=None, dtype=None)

import torch
from torch import nn


configs = [
    ("BatchNorm1d", nn.BatchNorm1d(5, momentum=0.3), (3, 5, 1)),
    ("BatchNorm2d", nn.BatchNorm2d(6, momentum=0.3), (3, 6, 2, 2)),
    ("BatchNorm3d", nn.BatchNorm3d(4, momentum=0.3), (3, 4, 2, 2, 3)),
]

for label, batch_norm, shape in configs:
    x = torch.ones(shape)
    for channel in range(shape[1]):
        x[:, channel] *= channel + 1  # チャンネルごとに値を変え、C軸だけがnum_featuresに対応することを確かめやすくする

    _ = batch_norm(x)
    print(
        label,
        "input",
        tuple(x.shape),
        "running_mean",
        tuple(batch_norm.running_mean.shape),
        "weight",
        tuple(batch_norm.weight.shape),
    )
    # BatchNorm1d input (3, 5, 1) running_mean (5,) weight (5,)
    # BatchNorm2d input (3, 6, 2, 2) running_mean (6,) weight (6,)
    # BatchNorm3d input (3, 4, 2, 2, 3) running_mean (4,) weight (4,)

import torch
from torch import nn
import torch.nn.functional as F


class LeNetBN(nn.Module):
    def __init__(self, classes):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, kernel_size=5)
        self.bn1 = nn.BatchNorm2d(num_features=6)

        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.bn2 = nn.BatchNorm2d(num_features=16)

        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 32x32入力を二回のconv+poolに通した後のflatten次元
        self.bn3 = nn.BatchNorm1d(num_features=120)

        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, classes)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)  # 畳み込み特徴 (N, C, H, W) にはBatchNorm2dを使う
        x = F.relu(x)
        x = F.max_pool2d(x, kernel_size=2)

        x = self.conv2(x)
        x = self.bn2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, kernel_size=2)

        x = x.view(x.size(0), -1)  # (N, C, H, W) -> (N, C*H*W) にflattenする

        x = self.fc1(x)
        x = self.bn3(x)  # flatten後の全結合特徴 (N, C) にはBatchNorm1dを使う
        x = F.relu(x)

        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = LeNetBN(classes=2)
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, gamma=0.1)

# train_loader、valid_loader、max_epochは、上で説明した自前のDatasetから用意するプレースホルダーです。
# このブロックは画像データセットに依存するため、単体では実行できません（構造を示す骨組みです）。
for epoch in range(max_epoch):
    net.train()
    for images, labels in train_loader:
        outputs = net(images)
        loss = criterion(outputs, labels)

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

    scheduler.step()

    net.eval()
    with torch.inference_mode():
        for images, labels in valid_loader:
            outputs = net(images)
            valid_loss = criterion(outputs, labels)

nn.LayerNorm(
    normalized_shape,
    eps=1e-5,
    elementwise_affine=True,
    bias=True,
    device=None,
    dtype=None,
)

norm = nn.LayerNorm(D)
y = norm(x)

import torch
from torch import nn


x = torch.tensor(
    [[1.0, 2.0, 3.0],
     [10.0, 20.0, 30.0]]
)
layer_norm = nn.LayerNorm(3, elementwise_affine=False)  # 学習可能なgamma/betaを使わず、標準化だけを見る
y = layer_norm(x)

print(torch.round(y, decimals=4))
# tensor([[-1.2247,  0.0000,  1.2247],
#         [-1.2247,  0.0000,  1.2247]])
print("row means:", torch.round(y.mean(dim=1), decimals=4))
# row means: tensor([0., 0.])
print("row vars:", torch.round(y.var(dim=1, unbiased=False), decimals=4))
# row vars: tensor([1., 1.])

import torch
from torch import nn


x = torch.ones(8, 6, 3, 4)

layer_norm = nn.LayerNorm([6, 3, 4])  # 入力の末尾3次元(C, H, W)=(6, 3, 4)と一致させる
y = layer_norm(x)
print("weight:", tuple(layer_norm.weight.shape), "output:", tuple(y.shape))
# weight: (6, 3, 4) output: (8, 6, 3, 4)

try:
    wrong_layer_norm = nn.LayerNorm([6, 3])  # わざと入力の末尾2次元(3, 4)と食い違う指定にする
    wrong_layer_norm(x)
except RuntimeError as error:
    print(type(error).__name__, str(error).splitlines()[0])
    # RuntimeError Given normalized_shape=[6, 3], expected input with shape [*, 6, 3], but got input of size[8, 6, 3, 4]

各 (n, c) について、H, W方向で mean / variance を計算する

nn.InstanceNorm2d(
    num_features,
    eps=1e-5,
    momentum=0.1,
    affine=False,
    track_running_stats=False,
    device=None,
    dtype=None,
)

import torch
from torch import nn


feature_map = torch.ones(2, 2)
feature_maps = torch.stack(
    [feature_map * (channel + 1) for channel in range(3)],
    dim=0,
)  # 3チャンネル、各チャンネルは定数値（1, 2, 3）だけを持つ
feature_maps_bs = torch.stack([feature_maps for _ in range(2)], dim=0)

instance_norm = nn.InstanceNorm2d(
    num_features=3,
    affine=False,
    track_running_stats=False,
)
outputs = instance_norm(feature_maps_bs)

print("input shape:", tuple(feature_maps_bs.shape))
# input shape: (2, 3, 2, 2)
print("output[0]:", outputs[0])
# output[0]: tensor([[[0., 0.],
#          [0., 0.]],
#
#         [[0., 0.],
#          [0., 0.]],
#
#         [[0., 0.],
#          [0., 0.]]])
print("channel means:", outputs.mean(dim=(0, 2, 3)))
# channel means: tensor([0., 0., 0.])

nn.GroupNorm(
    num_groups,
    num_channels,
    eps=1e-5,
    affine=True,
    device=None,
    dtype=None,
    *,
    bias=True,
)

num_groups = num_channels: InstanceNormに近い
num_groups = 1           : チャンネルと空間をまとめるLayerNormに近い

import torch
from torch import nn


try:
    group_norm = nn.GroupNorm(num_groups=3, num_channels=4)  # 4を3で割り切れない組み合わせ
    group_norm(torch.ones(2, 4, 2, 2))
except ValueError as error:
    print(type(error).__name__, str(error))
    # ValueError num_channels (4) must be divisible by num_groups (3)

import torch
from torch import nn


x = torch.arange(2 * 4 * 2 * 2, dtype=torch.float32).reshape(2, 4, 2, 2)

modules = [
    ("BN", nn.BatchNorm2d(4, affine=False, track_running_stats=False)),
    ("IN", nn.InstanceNorm2d(4, affine=False, track_running_stats=False)),
    ("GN", nn.GroupNorm(2, 4, affine=False)),
    ("LN", nn.LayerNorm((4, 2, 2), elementwise_affine=False)),
]

for label, module in modules:
    y = module(x)
    print(
        label,
        tuple(y.shape),
        round(y.mean().item(), 6),
        round(y.var(unbiased=False).item(), 6),
    )
    # BN (2, 4, 2, 2) -0.0 1.0
    # IN (2, 4, 2, 2) -0.0 0.999992
    # GN (2, 4, 2, 2) -0.0 0.999998
    # LN (2, 4, 2, 2) 0.0 1.0

nn.RMSNorm(
    normalized_shape,
    eps=None,
    elementwise_affine=True,
    device=None,
    dtype=None,
)

sync_model = nn.SyncBatchNorm.convert_sync_batchnorm(model)

典型的なCNN:
Conv -> BatchNorm -> ReLU

Pre-activation ResNet:
BatchNorm -> ReLU -> Conv

TransformerのPre-Norm:
x -> Norm -> Attention/MLP -> Dropout -> Residual Add

import torch
from torch import nn


class RegularizedCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),

            nn.Conv2d(64, 128, 3, padding=1, bias=False),
            nn.GroupNorm(8, 128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(p=0.3),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        x = self.features(x)   # (N, 3, 32, 32) -> (N, 128, 8, 8)
        x = self.pool(x)       # (N, 128, 8, 8) -> (N, 128, 1, 1)
        x = self.classifier(x) # (N, 128, 1, 1) -> (N, num_classes)
        return x


model = RegularizedCNN(num_classes=10)
example = torch.randn(2, 3, 32, 32)
print("output shape:", model(example).shape)
# output shape: torch.Size([2, 10])

no_decay = []
decay = []
for name, parameter in model.named_parameters():
    if parameter.ndim == 1 or name.endswith("bias"):
        no_decay.append(parameter)
    else:
        decay.append(parameter)

optimizer = torch.optim.AdamW(
    [
        {"params": decay, "weight_decay": 1e-2},
        {"params": no_decay, "weight_decay": 0.0},
    ],
    lr=1e-3,
)

print("decay count:", len(decay), "no_decay count:", len(no_decay))
# decay count: 3 no_decay count: 5

モデル容量が小さい                         モデル容量が大きい
Bias:    高い --------------------------> 低い
Variance:低い --------------------------> 高い
検証誤差: 高い ---- 下がる ---- 再び上がる

L1制約: ひし形に近い制約領域。角が座標軸に乗りやすく、0が生まれやすい。
L2制約: 円に近い制約領域。なめらかに縮めるが、0ぴったりにはなりにくい。

今の重み w_t
  -> (1 - eta * lambda) を掛けて少し縮める
  -> データlossの勾配で更新する
  -> 次の重み w_(t+1)

訓練前:  o--o--o--o
mask:    1  0  1  0
訓練時:  o--x--o--x   残った値を 1/(1-p) 倍
評価時:  o--o--o--o   そのまま通す

x
 -> batchや空間方向から mean / var を計算する
 -> x_hat = (x - mean) / sqrt(var + eps)
 -> y = gamma * x_hat + beta

入力: (N, C, H, W)

BN: 同じCについて、N/H/Wをまとめる
LN: 同じサンプルについて、C/H/Wをまとめる
IN: 同じサンプル・同じCについて、H/Wをまとめる
GN: 同じサンプルのチャンネルをG組に分け、各組とH/Wをまとめる