現実世界を動かす技術をシェアしよう！～身近なものづくりから最先端テクノロジーまで大募集！～

PyTorch入門第5回：重み初期化、損失関数、最適化手法で学習を安定させる

Last updated at 2026-08-19Posted at 2026-08-18

はじめに

第4回では、nn.Module を使ってモデルを作り、Conv2d、プーリング、Linear、活性化関数を組み合わせて画像分類用のCNNを組み立てました。

モデルを定義できるようになると、次に必要になるのは「どう学習させるか」です。学習では、入力をモデルに通し、出力と正解のずれを損失関数で数値化し、その損失が小さくなるように最適化手法でパラメータを更新します。

この記事では、重み初期化、損失関数、最適化手法をまとめて扱います。これらは別々の部品に見えますが、実際の訓練では強くつながっています。初期化が悪いと勾配が小さくなりすぎたり大きくなりすぎたりします。損失関数と出力の形が合っていないと、学習目標そのものがずれます。学習率や最適化手法の設定が合っていないと、lossが下がらなかったり発散したりします。

目的は、関数名を暗記することではありません。model(inputs) から loss.backward()、optimizer.step() までの流れを、自分で形、dtype、値の意味を確認しながら書けるようになることです。

この記事で扱うこと
第4回から第5回へのつながり
訓練ループの全体像
なぜ重み初期化が必要なのか
fan_in、fan_out、gain
Xavier初期化
Kaiming初期化
PyTorchの初期化関数
損失関数を読むときの基本
reduction の意味
nn.CrossEntropyLoss
nn.NLLLoss
BCEWithLogitsLoss と BCELoss
回帰で使う損失関数
nn.KLDivLoss
そのほかの損失関数の見方
損失関数の選び方
最適化手法の基本
zero_grad、backward、step
学習率
SGD、Momentum、Nesterov
Adam、AdamW、よく使う最適化手法
パラメータグループと凍結
勾配クリッピングとスケジューラ
小さな完全訓練ループ
チェックポイント保存と再開
よくあるつまずきどころ
参考リンク
おわりに

この記事で扱うこと

この記事では、次の内容を扱います。

勾配消失と勾配爆発の直感
fan_in、fan_out、gain の意味
Xavier初期化とKaiming初期化
torch.nn.init の主な関数
model.apply() を使った初期化
損失、目的関数、reduction の考え方
CrossEntropyLoss、NLLLoss、BCEWithLogitsLoss
L1、MSE、SmoothL1、Huber
KLDivLoss、ランキング系損失、CTCLossの位置づけ
optimizer.zero_grad()、loss.backward()、optimizer.step()
SGD、Momentum、Nesterov
Adam、AdamW、パラメータグループ
勾配クリッピング、学習率スケジューラ
チェックポイント保存と再開

コードは、手元に画像データがなくても動かせるように、乱数Tensorや小さなダミーデータを使います。実データに置き換えるときも、shape、dtype、損失関数の選び方は同じです。

第4回から第5回へのつながり

第4回で作った画像分類モデルは、入力画像を受け取り、各クラスに対応するlogitを返しました。

入力画像:  (N, 3, 224, 224)
モデル出力: (N, 2)
正解ラベル: (N,)

ここで重要なのは、モデルの出力が確率ではなくlogitであることです。logitはSoftmaxをかける前の生のスコアです。訓練で nn.CrossEntropyLoss を使う場合、モデルの最後でSoftmaxをかける必要はありません。CrossEntropyLoss が内部で数値的に安定した形で処理してくれます。

第4回の流れを訓練まで広げると、次のようになります。

Dataset / DataLoader
  -> images, labels
  -> model(images)
  -> logits
  -> criterion(logits, labels)
  -> loss
  -> loss.backward()
  -> optimizer.step()

この記事では、この後半部分を丁寧に見ます。モデルを作る力と、学習を安定して回す力は、セットで必要になります。

訓練ループの全体像

まず、訓練の1 stepを見ます。

import torch
from torch import nn


# 小さな分類モデルです。
# 入力は20次元、出力は3クラス分のlogitです。
model = nn.Sequential(
    nn.Linear(20, 32),  # 20次元の入力を32次元の隠れ表現へ変換する
    nn.ReLU(),          # 非線形性を入れる
    nn.Linear(32, 3),   # 3クラス分のlogitを出す
)

# CrossEntropyLossは、logitと整数のクラスラベルを受け取ります。
criterion = nn.CrossEntropyLoss()

# optimizerは、モデルのパラメータを更新する役割を持ちます。
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

# ダミーのmini-batchです。
inputs = torch.randn(8, 20)       # batch size 8、特徴量20次元
targets = torch.randint(0, 3, (8,))  # 各サンプルの正解クラス。dtypeはtorch.long

logits = model(inputs)            # 順伝播。shapeは（8, 3）
loss = criterion(logits, targets) # 予測と正解のずれを1つのスカラーにする

optimizer.zero_grad(set_to_none=True)  # 前回の勾配を消す
loss.backward()                        # lossから各パラメータの勾配を計算する
optimizer.step()                       # 勾配を使ってパラメータを更新する

print(logits.shape)
# torch.Size([8, 3])
print(loss.item())
# 0.9853639602661133

この短いコードの中に、訓練の重要部品がすべて入っています。

部品	役割
`model`	入力からlogitや予測値を作る
`criterion`	出力と正解のずれを損失として数値化する
`optimizer`	勾配を使ってパラメータを更新する
`zero_grad`	前回の勾配を消す
`backward`	autogradで勾配を計算する
`step`	optimizerがパラメータを1回更新する

この流れが正しく書けていても、初期化、損失関数、学習率が合っていないと学習は不安定になります。ここから順番に、その理由を見ます。

なぜ重み初期化が必要なのか

ニューラルネットワークでは、各層が入力に重みを掛け、次の層へ値を渡します。もし重みが大きすぎると、層を通るたびに値や勾配が大きくなり、lossが発散することがあります。逆に重みが小さすぎると、値や勾配が小さくなり、前の層がほとんど学習しないことがあります。

直感的には、次のような問題です。

重みが小さすぎる
  -> 層を通るたびに信号が小さくなる
  -> 勾配も小さくなる
  -> 前の層が学習しにくい

重みが大きすぎる
  -> 層を通るたびに信号が大きくなる
  -> 勾配も大きくなる
  -> lossがNaNやInfになりやすい

初期化の目的は、訓練開始時に、前向きの値と後ろ向きの勾配が極端に小さくも大きくもならないようにすることです。

次のコードでは、極端に大きい重みを入れた場合に、出力の値が大きくなりやすいことを確認します。

import torch
from torch import nn


torch.manual_seed(0)  # 結果を再現しやすくする

layer = nn.Linear(100, 100)  # 100次元から100次元へ変換する線形層
x = torch.randn(64, 100)     # batch size 64のダミー入力

with torch.no_grad():
    layer.weight.normal_(mean=0.0, std=5.0)  # あえて大きすぎる標準偏差で初期化する
    layer.bias.zero_()                       # biasは0にして重みの影響を見やすくする

y = layer(x)  # 大きい重みを通した出力

print(x.std())  # 入力の標準偏差
# tensor(1.0214)
print(y.std())  # 出力の標準偏差。かなり大きくなりやすい
# tensor(51.4013, grad_fn=<StdBackward0>)

次の例では、深いMLPを作り、各層を通った後の標準偏差を確認します。重みの標準偏差を大きくしすぎる場合と、Kaiming初期化を使う場合を比べると、初期化が信号の大きさに直接影響することが見えます。

import torch
from torch import nn


torch.manual_seed(1)


class DeepMLP(nn.Module):
    def __init__(self, width=128, depth=12):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(width, width, bias=False) for _ in range(depth)
        ])

    def forward(self, x):
        std_history = []
        for layer_index, layer in enumerate(self.layers):
            x = layer(x)                 # 線形変換で次の層へ渡す値を作る
            x = torch.relu(x)            # ReLUで負の値を0にする
            std_history.append(x.std())  # その層の出力スケールを記録する
        return x, std_history


def initialize_with_large_std(module):
    if isinstance(module, nn.Linear):
        nn.init.normal_(module.weight, mean=0.0, std=1.0)  # あえて大きすぎる標準偏差にする


def initialize_with_kaiming(module):
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, nonlinearity="relu")  # ReLU向けの初期化にする


inputs = torch.randn(16, 128)

for name, initializer in [
    ("std=1.0", initialize_with_large_std),
    ("kaiming", initialize_with_kaiming),
]:
    model = DeepMLP(width=128, depth=12)
    model.apply(initializer)

    _, std_history = model(inputs)
    print(name)
    print([round(value.item(), 4) for value in std_history[:4]])   # 最初の数層
    print([round(value.item(), 4) for value in std_history[-4:]])  # 最後の数層
    # 出力例:
    # std=1.0
    # [6.7331, 52.8879, 396.5564, 3459.4866]
    # [110020256.0, 873119936.0, 7058934784.0, 53469868032.0]
    # kaiming
    # [0.8536, 0.9098, 0.8443, 0.8032]
    # [0.7235, 0.6899, 0.5702, 0.5856]

std=1.0 のように大きい標準偏差で始めると、層を重ねるほど値のスケールが急に大きくなりやすくなります。一方、Kaiming初期化では、ReLUを通した後の値のスケールが極端に崩れにくくなります。

この標準偏差の急激な増加は、逆伝播時の勾配にも同じように表れます。値が大きくなりすぎると勾配も大きくなり、lossがNaNやInfになりやすくなります。最初に見た「重みが大きすぎる、小さすぎる」という問題が、層を重ねるほど実際に大きく効いてくることが、この例からも分かります。

この例は極端ですが、深いモデルでは小さなスケールのずれが層を重ねるほど大きく効いてきます。そのため、活性化関数に合った初期化を選ぶことが大切です。

fan_in、fan_out、gain

重み初期化を理解するとき、まず fan_in と fan_out を押さえます。

用語	意味	何を保ちたいか
`fan_in`	1つの出力ユニットに入ってくる接続数	前向きの値のスケール
`fan_out`	1つの入力ユニットから出ていく接続数	後ろ向きの勾配のスケール

Linear の重みは、PyTorchでは次の形です。

weight.shape = (out_features, in_features)

したがって、Linear(20, 32) なら fan_in=20、fan_out=32 です。

Conv2d の重みは、次の形です。

weight.shape = (out_channels, in_channels / groups, kernel_h, kernel_w)

たとえば、Conv2d(3, 16, kernel_size=3) では、1つの出力位置を見るために 3 * 3 * 3 = 27 個の入力値を使うため、fan_in は27になります。

gain は、活性化関数に合わせて初期化のスケールを補正する係数です。PyTorchでは nn.init.calculate_gain() で計算できます。

from torch import nn


relu_gain = nn.init.calculate_gain("relu")       # ReLU向けの補正係数
linear_gain = nn.init.calculate_gain("linear")   # 線形層向けの補正係数
tanh_gain = nn.init.calculate_gain("tanh")       # Tanh向けの補正係数

print(relu_gain)
# 1.4142135623730951
print(linear_gain)
# 1
print(tanh_gain)
# 1.6666666666666667

次のコードでは、乱数を tanh に通す前後で標準偏差を比べ、経験的なgainを計算します。同じ発想で見ると、calculate_gain("tanh") が「活性化関数で変わるスケールを補うための係数」だと理解しやすくなります。

import torch
from torch import nn


torch.manual_seed(1)

x = torch.randn(10000)  # 平均0、標準偏差1に近い乱数
out = torch.tanh(x)     # Tanhを通すと値が (-1, 1) に押し込まれる

empirical_gain = x.std() / out.std()            # 入力と出力の標準偏差の比を計算する
pytorch_gain = nn.init.calculate_gain("tanh")  # PyTorchが用意しているTanh向けgain

print(empirical_gain)
# tensor(1.5983)
print(pytorch_gain)
# 1.6666666666666667

この2つは完全に同じ値になる必要はありません。上の empirical_gain は、たまたま生成した乱数に対する実測値です。一方、calculate_gain("tanh") は、初期化で使いやすい代表的な係数として用意されています。

活性化関数が変わると、値の分布も変わります。たとえばReLUは負の値を0にするため、線形層と同じスケールの初期化では合わないことがあります。Kaiming初期化は、その点を考慮した初期化です。

Xavier初期化

Xavier初期化は、線形層やTanhのような対称的な活性化を想定して、前向きと後ろ向きのスケールを両方見ながら重みを決める考え方です。Glorot初期化とも呼ばれます。

基本の考え方は、重みの分散を次のように設定することです。

Var(W) = 2 / (fan_in + fan_out)

一様分布を使う場合は、次の範囲から値を取ります。

W ~ Uniform(-a, a)
a = gain * sqrt(6 / (fan_in + fan_out))

PyTorchでは次の関数を使います。

import torch
from torch import nn


linear = nn.Linear(20, 32)  # 20次元から32次元へ変換する層

with torch.no_grad():
    nn.init.xavier_uniform_(linear.weight)  # Xavierの一様分布で重みを初期化する
    nn.init.zeros_(linear.bias)             # biasを0にする

print(linear.weight.mean())
# 出力例（このブロックはmanual_seedを固定していないため、値は実行ごとに変わります）: tensor(-0.0041)
print(linear.weight.std())
# 出力例（同上）: tensor(0.2008)

Xavier初期化には、正規分布版もあります。

linear = nn.Linear(20, 32)

with torch.no_grad():
    nn.init.xavier_normal_(linear.weight)  # Xavierの正規分布で重みを初期化する
    nn.init.zeros_(linear.bias)

最初は、「Tanhや線形に近い層ではXavierを候補にする」と覚えると整理しやすいです。一方、ReLUを多く使うCNNやMLPでは、次のKaiming初期化がよく使われます。

Kaiming初期化

Kaiming初期化は、ReLUやLeakyReLUのように、負側の値を切る活性化関数に合わせた初期化です。He初期化とも呼ばれます。

ReLUでは、負の値が0になります。そのぶん出力の分布が変わるため、初期化のスケールも調整します。前向きの値のスケールを保つ考え方では、よく次の形で説明されます。

Var(W) ≈ 2 / fan_in
std(W) ≈ sqrt(2 / fan_in)

PyTorchでは次のように書きます。

import torch
from torch import nn


conv = nn.Conv2d(3, 32, kernel_size=3, padding=1, bias=False)

with torch.no_grad():
    nn.init.kaiming_normal_(
        conv.weight,          # 初期化したいTensor
        mode="fan_out",       # 勾配側のスケールを保つ設定としてよく使われる
        nonlinearity="relu",  # 後ろにReLUを置く前提でgainを決める
    )

print(conv.weight.shape)
# torch.Size([32, 3, 3, 3])
print(conv.weight.std())
# 出力例（このブロックはmanual_seedを固定していないため、値は実行ごとに変わります）: tensor(0.0805)

mode には主に2つあります。

`mode`	意味	よく考える方向
`fan_in`	入力側の接続数を基準にする	前向きの値のスケール
`fan_out`	出力側の接続数を基準にする	後ろ向きの勾配のスケール

CNNでは、Conv2dの後にReLUを置くことが多いため、Kaiming初期化が自然な候補になります。ただし、PyTorchの多くの層には初期化が最初から入っているため、すべてのモデルで必ず手動初期化が必要というわけではありません。自分で初期化を上書きするときは、モデルの構造と活性化関数を意識します。

PyTorchの初期化関数

torch.nn.init には、よく使う初期化関数がまとまっています。

関数	役割
`xavier_uniform_`	Xavier一様分布で初期化する
`xavier_normal_`	Xavier正規分布で初期化する
`kaiming_uniform_`	Kaiming一様分布で初期化する
`kaiming_normal_`	Kaiming正規分布で初期化する
`uniform_`	指定範囲の一様分布で初期化する
`normal_`	指定平均・標準偏差の正規分布で初期化する
`constant_`	指定した定数で埋める
`zeros_`	0で埋める
`ones_`	1で埋める
`orthogonal_`	直交行列に近い形で初期化する
`eye_`	2次元Tensorを単位行列のように初期化する
`dirac_`	畳み込みで入力をなるべくそのまま通す形にする
`sparse_`	2次元Tensorを疎な形で初期化する

関数名の最後に _ が付いている点にも注目します。PyTorchでは、末尾に _ が付く関数は、Tensorをその場で書き換える操作を表すことが多いです。

モデル全体に初期化を適用したい場合は、model.apply() が便利です。apply() は、モデル内のすべての子モジュールに関数を再帰的に適用します。

import torch
from torch import nn


class SmallMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(20, 64),
            nn.ReLU(),
            nn.Linear(64, 3),
        )

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


@torch.no_grad()
def initialize(module):
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, nonlinearity="relu")  # Linearの重みをKaimingで初期化する
        if module.bias is not None:
            nn.init.zeros_(module.bias)  # biasがある場合は0にする


model = SmallMLP()
model.apply(initialize)  # model内の各モジュールにinitializeを適用する

for name, parameter in model.named_parameters():
    print(name, parameter.shape)
    # net.0.weight torch.Size([64, 20])
    # net.0.bias torch.Size([64])
    # net.2.weight torch.Size([3, 64])
    # net.2.bias torch.Size([3])

ここで大切なのは、initialize() が nn.Linear だけを対象にしている点です。モデルには ReLU のように重みを持たない層も含まれます。isinstance() で層の種類を確認してから初期化すると、不要なエラーを避けられます。

CNNなら、Conv2d と Linear で初期化を分ける書き方もよくあります。

@torch.no_grad()
def initialize_cnn(module):
    if isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")  # Conv2dはReLU前提で初期化する
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Linear):
        nn.init.xavier_uniform_(module.weight)  # 最後の分類器はXavierで初期化する例
        if module.bias is not None:
            nn.init.zeros_(module.bias)

初期化は、訓練のスタート地点を決める操作です。学習済みの重みを読み込んだ後に全体を初期化し直すと、読み込んだ重みが消えてしまいます。初期化を呼ぶ場所にも注意します。

損失関数を読むときの基本

損失関数は、モデルの出力と正解のずれを数値にする関数です。

モデル出力 + 正解 -> loss

損失関数を選ぶときは、次の4つを必ず確認します。

確認するもの	例
モデル出力	logitか、確率か、連続値か
targetの形式	クラス番号か、0/1ラベルか、連続値か
dtype	`torch.long` か、浮動小数点か
shape	`(N, C)` か、`(N,)` か、同じ形か

ここを間違えると、コードは動いても学習の意味がずれることがあります。

たとえば、多クラス分類では、モデル出力はlogit、targetは整数のクラス番号にします。

logits:  (N, C)       # float
labels:  (N,)         # torch.long
loss:    scalar

二分類や多ラベル分類で BCEWithLogitsLoss を使う場合は、モデル出力とtargetを同じ形にします。

logits:  (N, 1) または (N, C)   # float
target:  同じ形                  # float、値は0または1
loss:    scalar

「どの損失関数を使うか」は、「モデルの最後をどう設計するか」と一体です。損失関数だけを見て決めるのではなく、出力層、target、評価指標までセットで考えます。

`reduction` の意味

多くの損失関数には reduction があります。これは、各サンプルや各要素の損失を、最終的にどうまとめるかを指定する引数です。

`reduction`	意味	返る値
`"none"`	各要素の損失をそのまま返す	Tensor
`"sum"`	損失をすべて足す	scalar
`"mean"`	損失を平均する	scalar

まず reduction="none" で形を確認すると、損失がどの単位で計算されているかを理解しやすくなります。

import torch
from torch import nn


logits = torch.tensor([
    [2.0, 0.1, -1.0],
    [0.0, 1.5, 0.2],
])
targets = torch.tensor([0, 2])

criterion_none = nn.CrossEntropyLoss(reduction="none")  # サンプルごとの損失を返す
criterion_mean = nn.CrossEntropyLoss(reduction="mean")  # batch内の損失を平均する

loss_each = criterion_none(logits, targets)
loss_mean = criterion_mean(logits, targets)

print(loss_each)
# tensor([0.1818, 1.7026])
print(loss_mean)
# tensor(0.9422)

reduction="none" は、サンプルごとの損失を見たいとき、重み付けを自分で行いたいとき、画像の各画素ごとの損失を可視化したいときに役立ちます。通常の訓練では、まず mean を使うことが多いです。

`nn.CrossEntropyLoss`

nn.CrossEntropyLoss は、単一ラベルの多クラス分類でよく使う損失関数です。画像分類で「1枚の画像は、犬・猫・鳥のどれか1つ」のように、正解クラスが1つだけの場合に使います。

基本形は次の通りです。

criterion = nn.CrossEntropyLoss(
    weight=None,          # クラスごとの重み。クラス不均衡があるときに使う
    ignore_index=-100,    # 指定したラベルを損失計算から除外する
    reduction="mean",     # none、sum、meanのいずれか
    label_smoothing=0.0,  # 正解ラベルを少しなめらかにする強度
)

CrossEntropyLoss は、モデル出力としてSoftmax後の確率ではなく、Softmax前のlogitを受け取ります。

import torch
from torch import nn


criterion = nn.CrossEntropyLoss()

logits = torch.tensor([
    [3.0, 0.5, -1.0],  # 1サンプル目はクラス0のスコアが高い
    [0.2, 0.1, 2.0],   # 2サンプル目はクラス2のスコアが高い
])
targets = torch.tensor([0, 2])  # 正解クラス番号。dtypeはtorch.long

loss = criterion(logits, targets)

print(loss)
# tensor(0.1847)

ここでSoftmaxを自分でかけないのが重要です。

実際のプロジェクトでは、紙幣の画像を2クラスに分類するデータセットを使うことがあります。画像を train と valid のフォルダに分け、DataLoader でbatchを作り、CNNモデルに通して CrossEntropyLoss を計算する、という流れです。この記事ではその画像データを用意できないため、同じ「2クラス分類」の形を小さなTensorで再現します。

import torch
from torch import nn


# 3サンプル、2クラス分のlogitです。
# 行ごとに1サンプル、列ごとにクラス0・クラス1のスコアを表します。
logits = torch.tensor([
    [1.0, 2.0],
    [1.0, 3.0],
    [1.0, 3.0],
])

# 正解クラス番号です。
# 1サンプル目はクラス0、2・3サンプル目はクラス1です。
targets = torch.tensor([0, 1, 1], dtype=torch.long)

loss_none = nn.CrossEntropyLoss(reduction="none")(logits, targets)  # サンプルごとの損失
loss_sum = nn.CrossEntropyLoss(reduction="sum")(logits, targets)    # 3サンプル分を合計
loss_mean = nn.CrossEntropyLoss(reduction="mean")(logits, targets)  # 3サンプル分を平均

print(loss_none)
# tensor([1.3133, 0.1269, 0.1269])
print(loss_sum)
# tensor(1.5671)
print(loss_mean)
# tensor(0.5224)

最初のサンプルだけを手計算すると、CrossEntropyLossが何をしているかが見えます。正解クラスのlogitを取り出し、全クラスの exp(logit) の合計を使って、次の形を計算しています。

loss = - 正解クラスのlogit + log(sum(exp(全クラスのlogit)))

import torch
from torch import nn


logits = torch.tensor([
    [1.0, 2.0],
    [1.0, 3.0],
    [1.0, 3.0],
])
targets = torch.tensor([0, 1, 1], dtype=torch.long)

sample_index = 0
sample_logits = logits[sample_index]             # 1サンプル目のlogit
target_class = targets[sample_index].item()      # 1サンプル目の正解クラス番号

correct_logit = sample_logits[target_class]      # 正解クラスのlogit
log_sum_exp = torch.log(torch.exp(sample_logits).sum())  # 全クラスのexpを足してlogを取る
manual_loss = -correct_logit + log_sum_exp       # CrossEntropyの1サンプル分

torch_loss = nn.CrossEntropyLoss(reduction="none")(logits, targets)[sample_index]

print(manual_loss)
# tensor(1.3133)
print(torch_loss)
# tensor(1.3133)

loss = criterion(logits, targets)  # 正しい。logitをそのまま渡す

CrossEntropyLoss は内部で、LogSoftmaxとNLLLossに相当する処理を安定した形で行います。そのため、次のように確率へ変換してから渡す必要はありません。

probabilities = torch.softmax(logits, dim=1)

上の probabilities は、表示や推論結果の解釈には便利です。しかし、訓練の損失にはlogitをそのまま渡します。

クラス番号target

通常の分類では、targetはクラス番号です。

logits.shape = (N, C)
target.shape = (N,)
target.dtype = torch.long

import torch
from torch import nn


torch.manual_seed(0)

criterion = nn.CrossEntropyLoss()
logits = torch.randn(8, 10)          # 8サンプル、10クラス分のlogit
targets = torch.randint(0, 10, (8,)) # 各サンプルの正解クラス番号

loss = criterion(logits, targets)

print(logits.shape)
# torch.Size([8, 10])
print(targets.shape)
# torch.Size([8])
print(targets.dtype)
# torch.int64
print(loss.item())
# 2.417919635772705（torch.manual_seed(0)によりこの通りに再現されます）

targetをone-hotにしない点が、初学者が特につまずきやすいところです。CrossEntropyLoss にクラス番号を渡す場合、targetのshapeは (N,) です。

画像の画素ごとの分類

セグメンテーションのように、各画素を分類する場合は、logitのshapeが (N, C, H, W)、targetのshapeが (N, H, W) になります。

import torch
from torch import nn


criterion = nn.CrossEntropyLoss()

logits = torch.randn(2, 4, 16, 16)      # 2枚、4クラス、16x16画素のlogit
targets = torch.randint(0, 4, (2, 16, 16))  # 各画素の正解クラス番号

loss = criterion(logits, targets)

print(loss.item())
# 出力例（乱数のシードを固定していないため、値は実行するたびに変わります）: 1.713318109512329

この場合も、targetにはクラス次元 C がありません。logit側だけがクラス次元を持ちます。

クラス重みと `ignore_index`

クラスの出現数に偏りがある場合は、weight でクラスごとの重みを指定できます。

import torch
from torch import nn


class_weight = torch.tensor([1.0, 2.0, 4.0])  # クラス2の損失を大きめに見る
criterion = nn.CrossEntropyLoss(weight=class_weight)

logits = torch.randn(5, 3)
targets = torch.tensor([0, 1, 2, 2, 1])

loss = criterion(logits, targets)

print(loss.item())
# 出力例（乱数シードを固定していないため、値は実行ごとに変わります）: 1.004454493522644

weight を指定すると、クラスごとの損失の重みが変わります。reduction="mean" のときは、単純にサンプル数で割るのではなく、該当サンプルのクラス重みを分母として使う点に注意します。

import torch
from torch import nn


logits = torch.tensor([
    [1.0, 2.0],
    [1.0, 3.0],
    [1.0, 3.0],
])
targets = torch.tensor([0, 1, 1], dtype=torch.long)
class_weight = torch.tensor([1.0, 2.0])  # クラス1の損失を2倍にする

loss_none = nn.CrossEntropyLoss(weight=class_weight, reduction="none")(logits, targets)
loss_mean = nn.CrossEntropyLoss(weight=class_weight, reduction="mean")(logits, targets)

target_weights = class_weight[targets]               # 各サンプルに対応する重み
manual_mean = loss_none.sum() / target_weights.sum()  # PyTorchのweighted meanと同じ分母

print(loss_none)
# tensor([1.3133, 0.2539, 0.2539])
print(loss_mean)
# tensor(0.3642)
print(manual_mean)
# tensor(0.3642)

ignore_index は、特定のラベルを損失計算から外したいときに使います。セグメンテーションで「ラベルなし画素」を無視したい場合などに便利です。

criterion = nn.CrossEntropyLoss(ignore_index=255)

logits = torch.randn(1, 3, 4, 4)
targets = torch.randint(0, 3, (1, 4, 4))
targets[0, 0, 0] = 255  # この位置だけ損失計算から外す

loss = criterion(logits, targets)

print(loss.item())
# 出力例（乱数シードを固定していないため、値は実行ごとに変わります）: 1.242615818977356

ignore_index は、クラス番号targetのときに使う設定です。targetを確率分布として渡す場合とは扱いが異なります。

ラベル平滑化

label_smoothing は、正解クラスだけを1.0にするのではなく、少しだけ他のクラスにも確率を分ける考え方です。モデルが正解クラスに過度に自信を持ちすぎるのを抑える目的で使われます。

criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

logits = torch.randn(8, 5)
targets = torch.randint(0, 5, (8,))

loss = criterion(logits, targets)

print(loss.item())
# 出力例（乱数シードを固定していないため、値は実行ごとに変わります）: 1.7539408206939697

最初は label_smoothing=0.0 のままで十分です。基本の分類が正しく動くようになってから、過学習対策や汎化性能の改善として検討します。

`nn.NLLLoss`

nn.NLLLoss は、log-probabilitiesを受け取る損失関数です。つまり、入力には torch.log_softmax() をかけた値を渡します。

import torch
from torch import nn


logits = torch.randn(8, 4)
targets = torch.randint(0, 4, (8,))

log_probs = torch.log_softmax(logits, dim=1)  # log-probabilitiesに変換する
criterion = nn.NLLLoss()
loss = criterion(log_probs, targets)

print(loss.item())
# 出力例（乱数シードを固定していないため、値は実行ごとに変わります）: 1.7070213556289673

CrossEntropyLoss は、実質的に log_softmax + NLLLoss の組み合わせとして理解できます。

CrossEntropyLoss(logits, target)
  ≒ NLLLoss(log_softmax(logits), target)

そのため、手元にlogitがある普通の分類では、CrossEntropyLoss を直接使うほうが書きやすいです。NLLLoss は、自分でlog-probabilitiesを明示的に扱いたいときに使います。

`BCEWithLogitsLoss` と `BCELoss`

二分類や多ラベル分類では、各クラスを独立した0/1判定として扱うことがあります。このときよく使うのが nn.BCEWithLogitsLoss です。

BCEWithLogitsLoss は、Sigmoidと二値交差エントロピーをまとめた損失関数です。入力にはSigmoid前のlogitを渡します。

import torch
from torch import nn


criterion = nn.BCEWithLogitsLoss()

logits = torch.tensor([[1.2], [-0.7], [0.3]])  # Sigmoid前のスコア
targets = torch.tensor([[1.0], [0.0], [1.0]])  # 0または1のfloatラベル

loss = criterion(logits, targets)
probabilities = torch.sigmoid(logits)          # 表示や判定に使う確率

print(loss.item())
# 0.4069412648677826
print(probabilities)
# tensor([[0.7685],
#         [0.3318],
#         [0.5744]])

CrossEntropyLoss と違い、targetは浮動小数点です。0 と 1 の値を持ちますが、dtypeは torch.float32 のような浮動小数点にします。

多ラベル分類では、1つのサンプルに複数のラベルが同時に立つことがあります。

import torch
from torch import nn


criterion = nn.BCEWithLogitsLoss()

logits = torch.randn(4, 5)  # 4サンプル、5ラベル分のlogit
targets = torch.tensor([
    [1.0, 0.0, 0.0, 1.0, 0.0],
    [0.0, 1.0, 0.0, 0.0, 0.0],
    [1.0, 1.0, 0.0, 0.0, 1.0],
    [0.0, 0.0, 0.0, 1.0, 1.0],
])

loss = criterion(logits, targets)

print(loss.item())
# 出力例（乱数シードを固定していないため、値は実行ごとに変わります）: 0.7332664132118225

多クラス分類と多ラベル分類の違いは、ここでとても重要です。

タスク	正解の考え方	出力	よく使う損失
多クラス分類	1つだけ正解	`(N, C)` logits	`CrossEntropyLoss`
多ラベル分類	複数同時に正解	`(N, C)` logits	`BCEWithLogitsLoss`

BCELoss も二値交差エントロピーですが、入力に確率を渡す必要があります。

probabilities = torch.sigmoid(logits)
loss = nn.BCELoss()(probabilities, targets)

実用上は、Sigmoidと損失をまとめた BCEWithLogitsLoss のほうが数値的に扱いやすく、コードも間違えにくいです。まずは BCEWithLogitsLoss を基本として覚えるとよいです。

`pos_weight`

正例が少ない二分類や多ラベル分類では、pos_weight を使って正例側の損失を大きくできます。

import torch
from torch import nn


pos_weight = torch.tensor([3.0])  # 正例の損失を3倍にする
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)

logits = torch.randn(6, 1)
targets = torch.tensor([[1.0], [0.0], [0.0], [0.0], [1.0], [0.0]])

loss = criterion(logits, targets)

print(loss.item())
# 出力例（乱数シードを固定していないため、値は実行ごとに変わります）: 0.944580078125

次の例では、BCELoss と BCEWithLogitsLoss を同じ入力で比べます。BCELoss では入力を先にSigmoidで確率へ変換し、BCEWithLogitsLoss ではlogitをそのまま渡します。

import torch
from torch import nn


logits = torch.tensor([
    [1.0, 2.0],
    [2.0, 2.0],
    [3.0, 4.0],
    [4.0, 5.0],
])
targets = torch.tensor([
    [1.0, 0.0],
    [1.0, 0.0],
    [0.0, 1.0],
    [0.0, 1.0],
])

probabilities = torch.sigmoid(logits)  # BCELossへ渡すために確率へ変換する

bce_loss = nn.BCELoss(reduction="none")(probabilities, targets)
bce_with_logits_loss = nn.BCEWithLogitsLoss(reduction="none")(logits, targets)

print(bce_loss)
# tensor([[0.3133, 2.1269],
#         [0.1269, 2.1269],
#         [3.0486, 0.0181],
#         [4.0181, 0.0067]])
print(bce_with_logits_loss)
# tensor([[0.3133, 2.1269],
#         [0.1269, 2.1269],
#         [3.0486, 0.0181],
#         [4.0181, 0.0067]])

2つの結果はほぼ同じ意味を持ちますが、実際の訓練では BCEWithLogitsLoss を優先すると、Sigmoidとlog計算をまとめて数値的に安定した形で扱えます。

pos_weight は、分類しきい値そのものを変える設定ではありません。損失の中で正例の重みを変える設定です。最終的なPrecisionやRecallは、検証データでしきい値も含めて確認します。

回帰で使う損失関数

連続値を予測する回帰では、分類用の交差エントロピーではなく、予測値と正解値の差を直接見る損失を使います。

`nn.MSELoss`

MSELoss は、差を二乗します。大きな誤差を強く罰するため、外れ値の影響を受けやすいです。

import torch
from torch import nn


criterion = nn.MSELoss()

prediction = torch.tensor([[2.5], [0.0], [4.0]])
target = torch.tensor([[3.0], [-1.0], [2.0]])

loss = criterion(prediction, target)

print(loss.item())
# 1.75

`nn.L1Loss`

L1Loss は、差の絶対値を使います。大きな外れ値に対して、MSEよりも影響を抑えやすいです。

criterion = nn.L1Loss()

prediction = torch.tensor([[2.5], [0.0], [4.0]])
target = torch.tensor([[3.0], [-1.0], [2.0]])

loss = criterion(prediction, target)

print(loss.item())
# 1.1666666269302368

`nn.SmoothL1Loss` と `nn.HuberLoss`

SmoothL1Loss と HuberLoss は、小さい誤差では二乗誤差のように振る舞い、大きい誤差では絶対値誤差のように振る舞います。外れ値にある程度強く、かつ0付近では滑らかに扱えます。

import torch
from torch import nn


prediction = torch.tensor([[2.5], [0.0], [4.0]])
target = torch.tensor([[3.0], [-1.0], [2.0]])

smooth_l1 = nn.SmoothL1Loss(beta=1.0)
huber = nn.HuberLoss(delta=1.0)

loss_smooth_l1 = smooth_l1(prediction, target)
loss_huber = huber(prediction, target)

print(loss_smooth_l1.item())
# 0.7083333134651184
print(loss_huber.item())
# 0.7083333134651184

回帰では、評価指標としてMAEやRMSEを別に見ることも多いです。訓練で使う損失と、最終的に報告する指標は、同じでなくてもかまいません。

`nn.KLDivLoss`

KLDivLoss は、分布どうしの差を扱う損失です。知識蒸留や、モデルの出力分布を別の分布に近づけたい場面で使われます。

PyTorchの KLDivLoss では、通常、入力にはlog-probabilitiesを渡し、targetにはprobabilitiesを渡します。

import torch
from torch import nn


student_logits = torch.tensor([[2.0, 0.5, -1.0]])
teacher_logits = torch.tensor([[1.5, 1.0, -0.5]])

student_log_probs = torch.log_softmax(student_logits, dim=1)  # 入力はlog-probabilities
teacher_probs = torch.softmax(teacher_logits, dim=1)          # targetはprobabilities

criterion = nn.KLDivLoss(reduction="batchmean")
loss = criterion(student_log_probs, teacher_probs)

print(loss.item())
# 出力例（この例は固定値だが、別の値で試すと結果は変わります）: 0.11225734651088715

reduction="batchmean" は、KLの数学的な意味に近い形で扱いたいときによく使います。mean と batchmean は同じではないため、KLを使うときは reduction を意識します。

交差エントロピー、エントロピー、KL散度の関係は、次のように整理できます。

交差エントロピー = エントロピー + KL散度

正解分布が固定されている場合、交差エントロピーを小さくすることは、予測分布を正解分布に近づけることとつながります。分類で CrossEntropyLoss がよく使われる背景にも、この考え方があります。

そのほかの損失関数の見方

PyTorchには、多くの損失関数があります。すべてを一度に暗記する必要はありません。まずは、「何を予測しているか」「targetがどんな形か」で分類すると見通しがよくなります。

損失関数	主な用途	まず確認すること
`PoissonNLLLoss`	カウントデータ	入力がlog-rateかrateか
`GaussianNLLLoss`	平均と分散を予測する回帰	分散が正の値になっているか
`MarginRankingLoss`	2つのスコアの順位づけ	targetが1または-1か
`TripletMarginLoss`	距離学習	anchor、positive、negativeの形
`CosineEmbeddingLoss`	ベクトルの類似・非類似	targetが1または-1か
`CTCLoss`	音声認識やOCRなどの未整列系列	入力が `(T, N, C)` のlog-probabilitiesか

CTCのような専用損失は、入力shapeの約束が特に重要です。通常の分類と同じ感覚で (N, T, C) を渡すと合いません。損失関数ごとに「期待する入力」と「期待するtarget」を必ず確認します。

ここでは、代表的なものを小さなダミーTensorでまとめて確認します。

import torch
from torch import nn


# MarginRankingLoss: target=1ならx1をx2より大きくしたい、target=-1なら小さくしたい
x1 = torch.tensor([[1.0], [2.0], [3.0]])
x2 = torch.tensor([[2.0], [2.0], [2.0]])
ranking_target = torch.tensor([[1.0], [1.0], [-1.0]])
ranking_loss = nn.MarginRankingLoss(margin=0.0, reduction="none")(x1, x2, ranking_target)

# TripletMarginLoss: anchorはpositiveに近く、negativeから遠くなるようにする
anchor = torch.tensor([[1.0]])
positive = torch.tensor([[2.0]])
negative = torch.tensor([[0.5]])
triplet_loss = nn.TripletMarginLoss(margin=1.0, p=1)(anchor, positive, negative)

# CTCLoss: 入力は (T, N, C) のlog-probabilitiesにする
T, N, C = 8, 2, 5
log_probs = torch.randn(T, N, C).log_softmax(dim=2).detach().requires_grad_()
ctc_targets = torch.tensor([[1, 2, 1], [1, 3, 0]], dtype=torch.long)
input_lengths = torch.full((N,), T, dtype=torch.long)
target_lengths = torch.tensor([3, 2], dtype=torch.long)
ctc_loss = nn.CTCLoss(blank=0)(log_probs, ctc_targets, input_lengths, target_lengths)

print(ranking_loss)
# tensor([[1.],
#         [0.],
#         [1.]])
print(triplet_loss)
# tensor(1.5000)
print(ctc_loss)
# 出力例（log_probsに乱数を使っていてseedを固定していないため、値は実行ごとに変わります）: tensor(3.5175, grad_fn=<MeanBackward0>)

このコードの目的は、各損失を実務ですぐ使い分けることではなく、損失関数ごとに入力とtargetの約束が大きく違うことを確認することです。実際に使うときは、公式ドキュメントで入力とtargetの形を確認したうえで、まず小さなダミーTensorで一度動かしてみると安心です。

損失関数の選び方

代表的な選び方をまとめると、次のようになります。

タスク	モデル出力	target	損失関数
単一ラベル多クラス分類	logits `(N, C)`	`long` クラス番号 `(N,)`	`CrossEntropyLoss`
二分類	logits `(N, 1)` または `(N,)`	float 0/1	`BCEWithLogitsLoss`
多ラベル分類	logits `(N, C)`	float multi-hot `(N, C)`	`BCEWithLogitsLoss`
通常の回帰	連続値	連続値	`MSELoss`、`L1Loss`
外れ値に強い回帰	連続値	連続値	`SmoothL1Loss`、`HuberLoss`
分布の近似	log-probabilities	probabilities	`KLDivLoss`
距離学習	埋め込みベクトル	ペアや三つ組	Triplet系、Cosine系
未整列系列	`(T, N, C)` log-probabilities	ラベル列と長さ	`CTCLoss`

損失選びで迷ったときは、次の順番で考えます。

予測したいものはクラスか、連続値か、分布か
1サンプルに正解クラスはいくつあるか
モデルの最後はlogitを出すのか、確率を出すのか
targetのshapeとdtypeは損失関数の期待と合っているか
reduction は訓練目的に合っているか

最適化手法の基本

最適化手法は、パラメータの .grad を見て、パラメータを更新します。

loss.backward()
  -> parameter.grad に勾配が入る
optimizer.step()
  -> parameter と optimizer内部状態が更新される

optimizerは、勾配を自分で作るわけではありません。勾配を作るのは loss.backward() です。optimizerは、その勾配を読んで更新します。

import torch
from torch import nn


model = nn.Linear(4, 2)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

inputs = torch.randn(3, 4)
targets = torch.tensor([0, 1, 0])

logits = model(inputs)
loss = criterion(logits, targets)

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

for name, parameter in model.named_parameters():
    print(name, parameter.grad.shape)  # backward後はgradが入っている
    # weight torch.Size([2, 4])
    # bias torch.Size([2])

optimizer.step()

optimizerは、アルゴリズムごとに内部状態を持つことがあります。たとえばSGDのmomentum buffer、AdamWの一階モーメント・二階モーメントなどです。この内部状態があるため、訓練を再開するときはモデルだけでなくoptimizerの状態も保存する必要があります。

次の例では、step()、zero_grad()、add_param_group()、state_dict() の働きを1つずつ確認します。小さなTensorだけでも、optimizerが「Tensorそのものへの参照」と「内部状態」を持っていることを確認できます。

import torch


weight = torch.randn(2, 2, requires_grad=True)
weight.grad = torch.ones(2, 2)  # 勾配がすべて1だと仮定する

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

print("before step")
print(weight.data)
# 出力例（weightがtorch.randnで乱数初期化されているため、値は実行するたびに変わります）:
# tensor([[-1.1626, -1.0768],
#         [ 0.3076, -3.0110]])

optimizer.step()  # weight = weight - lr * grad に近い更新が行われる

print("after step")
print(weight.data)
# 出力例（上と同じ理由で値は変わります）:
# tensor([[-1.2626, -1.1768],
#         [ 0.2076, -3.1110]])

print("grad before zero_grad")
print(weight.grad)
# tensor([[1., 1.],
#         [1., 1.]])

optimizer.zero_grad(set_to_none=True)  # gradをNoneに戻す

print("grad after zero_grad")
print(weight.grad)
# None

new_weight = torch.randn(3, 3, requires_grad=True)
optimizer.add_param_group({"params": [new_weight], "lr": 1e-4})  # 後からパラメータグループを追加する

print(len(optimizer.param_groups))
# 2
print(optimizer.state_dict().keys())
# dict_keys(['state', 'param_groups'])

optimizer.state_dict() には、パラメータグループとoptimizer内部状態が入ります。momentumつきSGDならmomentum buffer、AdamWなら移動平均の状態がここに保存されます。add_param_group() は、訓練の途中で新しいパラメータを追加したいときに使う方法です。最初から複数のパラメータグループに分けて指定する方法は、後述の「パラメータグループと凍結」で見ます。

`zero_grad`、`backward`、`step`

訓練ループで特に重要なのが、次の3行です。

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

それぞれの役割は明確に違います。

メソッド	役割
`zero_grad`	前回の勾配を消す
`backward`	lossから各パラメータの勾配を計算する
`step`	optimizerがパラメータを更新する

PyTorchでは、勾配は自動的に上書きされるのではなく、基本的に加算されます。そのため、通常のmini-batch訓練では、各stepの前に勾配を消します。

import torch
from torch import nn


model = nn.Linear(5, 2)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

for step in range(3):
    inputs = torch.randn(4, 5)       # ダミー入力
    targets = torch.randint(0, 2, (4,))  # ダミーラベル

    logits = model(inputs)           # 順伝播
    loss = criterion(logits, targets) # 損失を計算する

    optimizer.zero_grad(set_to_none=True)  # このstep用に勾配を空にする
    loss.backward()                        # 勾配を計算する
    optimizer.step()                       # パラメータを更新する

    print(step, loss.item())
    # 出力例（毎回inputs/targetsを乱数で作っているため、値は実行ごとに変わります）:
    # 0 0.8623557686805725
    # 1 0.6955071091651917
    # 2 1.0077991485595703

set_to_none=True は、勾配Tensorを0で埋めるのではなく None にします。多くの場合、メモリや速度の面で有利です。普通の訓練ではこの設定で問題ありません。

学習率

学習率は、1回の更新でどれくらいパラメータを動かすかを決める値です。

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

学習率が大きすぎると、最小値を飛び越えて発散しやすくなります。小さすぎると、lossがほとんど下がらず、訓練に時間がかかります。

簡単な1次元の例で見ます。

import torch


x = torch.tensor(2.0, requires_grad=True)
lr = 0.1

for step in range(5):
    loss = 4 * x ** 2           # 最小値はx=0
    loss.backward()             # dloss/dx = 8x

    with torch.no_grad():
        x -= lr * x.grad        # 勾配降下でxを更新する
        x.grad = None           # 次のstepのために勾配を消す

    print(step, x.item(), loss.item())
    # 0 0.4 16.0
    # 1 0.07999998331069946 0.6399999260902405
    # 2 0.015999995172023773 0.025599990040063858
    # 3 0.0031999992206692696 0.0010239994153380394
    # 4 0.0006399997510015965 4.0959981561172754e-05

この例では、lr を大きくしすぎると、値が振動したり発散したりします。ニューラルネットワークでも同じで、lossが急にNaNになったり大きく振動したりする場合は、学習率を最初に疑う価値があります。

次の例では、複数の学習率を並べて、同じ関数 y = 4x^2 をどのように下るかを比較します。グラフの代わりに、数step後のlossを見ます。

import torch


def function(x):
    return 4 * x ** 2  # 最小値はx=0


for lr in [0.01, 0.1, 0.2, 0.3]:
    x = torch.tensor(2.0, requires_grad=True)
    losses = []

    for step in range(8):
        loss = function(x)
        loss.backward()

        with torch.no_grad():
            x -= lr * x.grad
            x.grad = None

        losses.append(round(loss.item(), 4))

    print(f"lr={lr}", losses[-3:])
    # lr=0.01 [6.9502, 5.8827, 4.9791]
    # lr=0.1 [0.0, 0.0, 0.0]
    # lr=0.2 [0.0967, 0.0348, 0.0125]
    # lr=0.3 [462.8077, 907.1033, 1777.9227]

lr=0.01 はゆっくり進み、lr=0.1 は速く0に近づきます。lr=0.2 になると値が正負に振動しながら少しずつ0へ近づき、lr=0.3 のように大きくしすぎると、この単純な二次関数でも値が増え続けてしまいます。

SGD、Momentum、Nesterov

torch.optim.SGD は、基本的な勾配降下に近いoptimizerです。momentumを加えると、これまでの勾配方向をある程度覚えながら進みます。

import torch
from torch import nn


model = nn.Linear(10, 3)

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.01,        # 学習率
    momentum=0.9,   # これまでの勾配方向をどれくらい残すか
    weight_decay=1e-4,  # 重みを小さく保つための係数
    nesterov=True,  # Nesterov momentumを使う
)

print(optimizer)
# SGD (
# Parameter Group 0
#     dampening: 0
#     differentiable: False
#     foreach: None
#     fused: None
#     lr: 0.01
#     maximize: False
#     momentum: 0.9
#     nesterov: True
#     weight_decay: 0.0001
# )

Momentumの直感は、坂道を下るときの勢いです。同じ方向の勾配が続くと、その方向に進みやすくなります。一方で、方向が毎回大きく変わる場合は、振動を少し抑えられます。

勾配だけの更新:
  そのstepの勾配だけを見る

momentumつきの更新:
    これまでの勾配方向も少し覚えて進む

Nesterovは、少し先を見た位置で勾配を考えるような更新です。細かい式を最初から暗記する必要はありません。まずは、SGDを使うときに momentum=0.9 を付けることが多く、必要に応じて nesterov=True も検討される、と理解しておけば十分です。

次の例では、同じ関数に対して、momentumなしのSGDとmomentumつきSGDを比べます。最終的なlossだけを表示し、momentumが更新の進み方を変えることを確認します。

import torch


def function(x):
    return 4 * x ** 2


for momentum in [0.0, 0.9]:
    x = torch.tensor([2.0], requires_grad=True)
    optimizer = torch.optim.SGD([x], lr=0.03, momentum=momentum)
    losses = []

    for step in range(20):
        loss = function(x)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        losses.append(loss.item())

    print(f"momentum={momentum}", round(losses[-1], 6), round(x.item(), 6))
    # momentum=0.0 0.000473 0.008266
    # momentum=0.9 1.976168 -0.471941

momentumを入れると、更新は単に「その場の勾配だけ」で決まらなくなります。この例では、同じ lr=0.03 のままmomentumだけ0.9にすると、更新の勢いが強くなりすぎて、かえって振動が大きく残る結果になっています。momentumを上げるときは、lr を少し下げてバランスを取ることがよくあります。momentumは便利な設定ですが、学習率とセットで調整します。

Adam、AdamW、よく使う最適化手法

Adamは、勾配の平均と二乗平均を使って、パラメータごとに更新量を調整するoptimizerです。AdamWは、weight decayの扱いを分けたoptimizerで、多くのモデルで使われます。

import torch
from torch import nn


model = nn.Linear(10, 3)

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-3,          # AdamWでよく使われる開始点の一例
    betas=(0.9, 0.999),  # 一階・二階モーメントの移動平均係数
    eps=1e-8,         # 数値安定化のための小さい値
    weight_decay=1e-2,  # 解きたい問題に合わせて調整する
)

print(optimizer)
# AdamW (
# Parameter Group 0
#     amsgrad: False
#     betas: (0.9, 0.999)
#     capturable: False
#     decoupled_weight_decay: True
#     differentiable: False
#     eps: 1e-08
#     foreach: None
#     fused: None
#     lr: 0.001
#     maximize: False
#     weight_decay: 0.01
# )

optimizerの選び方に絶対の正解はありません。最初の目安としては、次のように考えると始めやすいです。

optimizer	特徴	よくある使い方
SGD	シンプルで、momentumと組み合わせやすい	CNNをじっくり訓練する場合
Adam	更新量をパラメータごとに調整する	早く動く基準モデルを作る場合
AdamW	weight decayを扱いやすい	画像モデル、Transformer、微調整
RMSprop	二乗勾配の移動平均を使う	一部の時系列・強化学習
SparseAdam	sparse gradient向け	Embeddingなど
LBFGS	closureを使う準ニュートン法	小規模で滑らかな問題

名前よりも、学習率、weight decay、スケジューラ、batch size、データ前処理のほうが結果に強く効くことも多いです。optimizerを変える前に、lossの形、target、学習率を確認します。

パラメータグループと凍結

optimizerには、パラメータグループを渡せます。これにより、層ごとに学習率やweight decayを変えられます。

import torch
from torch import nn


class TinyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Linear(20, 64),
            nn.ReLU(),
        )
        self.classifier = nn.Linear(64, 3)

    def forward(self, x):
        x = self.backbone(x)
        return self.classifier(x)


model = TinyModel()

optimizer = torch.optim.AdamW([
    {"params": model.backbone.parameters(), "lr": 1e-4},   # backboneは小さめの学習率
    {"params": model.classifier.parameters(), "lr": 1e-3}, # classifierは大きめの学習率
], weight_decay=1e-2)

print(len(optimizer.param_groups))
# 2

一部の層を更新したくない場合は、requires_grad_(False) を使って凍結します。

for parameter in model.backbone.parameters():
    parameter.requires_grad_(False)  # backboneのパラメータは勾配を計算しない

optimizer = torch.optim.AdamW(
    model.classifier.parameters(),  # 更新対象はclassifierだけ
    lr=1e-3,
)

凍結したパラメータをoptimizerに入れたままにするより、更新したいパラメータだけを渡すほうが見通しがよくなります。あとから凍結を解除する場合は、optimizerのパラメータグループも見直します。

biasや正規化層のパラメータにはweight decayをかけない、という分け方もよく使われます。

decay = []
no_decay = []

for name, parameter in model.named_parameters():
    if not parameter.requires_grad:
        continue
    if parameter.ndim == 1 or name.endswith("bias"):
        no_decay.append(parameter)  # biasや正規化層の重みはweight decayなしにする
    else:
        decay.append(parameter)     # それ以外の重みにはweight decayをかける

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

print(len(decay), len(no_decay))
# 1 1

この分け方はよく使われますが、すべてのモデルで必ず正解というわけではありません。まずは意図を理解したうえで使います。

勾配クリッピングとスケジューラ

勾配クリッピングは、勾配が大きくなりすぎたときに更新を抑えるための保護策です。特に長い系列モデルや、lossが不安定になりやすい設定で使われます。

import torch
from torch import nn


model = nn.Linear(10, 2)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

inputs = torch.randn(8, 10)
targets = torch.randint(0, 2, (8,))

logits = model(inputs)
loss = criterion(logits, targets)

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

total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 勾配ノルムを1.0以下に抑える
optimizer.step()

print(total_norm)
# 出力例（inputs/targetsを乱数で作っていてseedを固定していないため、値は実行ごとに変わります）: tensor(0.6635)

勾配クリッピングは、根本原因を直すものではありません。学習率が大きすぎる、入力のスケールが変、損失関数と出力が合っていない、といった問題がある場合は、そこも確認します。

学習率スケジューラは、訓練中に学習率を変える仕組みです。

import torch
from torch import nn


model = nn.Linear(10, 2)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer,
    step_size=5,  # 5 epochごとに学習率を下げる
    gamma=0.1,    # 学習率に0.1を掛ける
)

for epoch in range(10):
    inputs = torch.randn(8, 10)
    targets = torch.randint(0, 2, (8,))

    logits = model(inputs)
    loss = criterion(logits, targets)

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

    scheduler.step()  # optimizer.step() の後に学習率を更新する
    print(epoch, scheduler.get_last_lr())
    # 0 [0.001]
    # 1 [0.001]
    # 2 [0.001]
    # 3 [0.001]
    # 4 [0.0001]
    # 5 [0.0001]
    # 6 [0.0001]
    # 7 [0.0001]
    # 8 [0.0001]
    # 9 [1e-05]

スケジューラには、epochごとに呼ぶもの、batchごとに呼ぶもの、検証指標を見て呼ぶものがあります。使うスケジューラごとに、どのタイミングで step() するかを確認します。

小さな完全訓練ループ

ここまでの内容をまとめて、ダミーデータで動く小さな訓練ループを書きます。

実際のプロジェクトでの完全な訓練ループでは、紙幣画像の2クラス分類データセットを使うことがあります。画像を train と valid に分け、Resize、RandomCrop、RandomGrayscale、Normalize などの前処理をかけ、CNNモデルを CrossEntropyLoss とSGDで訓練する構成です。この記事ではその画像データを用意できないため、同じ5ステップの形だけを、乱数Tensorの分類問題として再現します。

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


torch.manual_seed(42)

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

# ダミーデータを作ります。
# 20次元の特徴量から3クラスを分類する練習用データです。
features = torch.randn(256, 20)
labels = torch.randint(0, 3, (256,))

train_loader = DataLoader(
    TensorDataset(features, labels),
    batch_size=32,
    shuffle=True,
)


class SmallClassifier(nn.Module):
    def __init__(self, input_dim=20, hidden_dim=64, num_classes=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_classes),
        )

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


@torch.no_grad()
def initialize(module):
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, nonlinearity="relu")  # ReLUを使うMLPとして初期化する
        if module.bias is not None:
            nn.init.zeros_(module.bias)


model = SmallClassifier().to(device)
model.apply(initialize)

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

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

    for inputs, targets in train_loader:
        inputs = inputs.to(device)
        targets = targets.to(device)

        logits = model(inputs)            # shapeは（batch, 3）
        loss = criterion(logits, targets) # targetsはshape（batch）のクラス番号

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

        running_loss += loss.item() * inputs.size(0)

    epoch_loss = running_loss / len(train_loader.dataset)
    print(f"epoch={epoch}, loss={epoch_loss:.4f}")
    # 出力例（torch.manual_seed(42)によりこの通りに再現されます）:
    # epoch=0, loss=1.6166
    # epoch=1, loss=1.4828
    # epoch=2, loss=1.3891

このデータは乱数なので、精度そのものには意味がありません。ここで確認したいのは、訓練ループの部品が正しい順番でつながっていることです。

実データに置き換えるときは、次を確認します。

features と labels を実データのTensorやDatasetに置き換える
モデルの入力次元とデータのshapeを合わせる
分類ならtargetを torch.long のクラス番号にする
二分類・多ラベルなら損失関数を BCEWithLogitsLoss にする
deviceをモデルと入力でそろえる

チェックポイント保存と再開

訓練を途中から再開したい場合は、モデルだけでなくoptimizerの状態も保存します。optimizerにはmomentumやAdamWの内部状態が入っているためです。

import torch


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

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

読み込むときは、同じ構造のモデルとoptimizerを作ってから load_state_dict() します。

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

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

print(start_epoch)
# 4

モデル重みだけを読み込むと、推論や微調整の開始には使えます。しかし、訓練を完全に続きから再開したい場合は、optimizerやschedulerの状態も必要です。学習率スケジューラや混合精度のscalerを使っている場合は、それらも一緒に保存します。

よくあるつまずきどころ

つまずき	原因	対処
`CrossEntropyLoss` にSoftmax後の確率を渡す	logitを期待している損失に確率を渡している	モデル出力のlogitをそのまま渡す
`CrossEntropyLoss` のtargetをone-hotにする	クラス番号targetの形式と混同している	targetは `(N,)` の `torch.long` にする
`BCEWithLogitsLoss` のtargetが整数型	BCE系は浮動小数点targetを期待する	0/1のfloat Tensorにする
多クラス分類と多ラベル分類を混同する	正解が1つか複数かを確認していない	多クラスはCE、多ラベルはBCEWithLogitsを基本にする
`zero_grad()` を忘れる	勾配が前stepから累積する	各stepで `zero_grad` を呼ぶ
`loss.backward()` 前に `optimizer.step()` する	勾配がまだ計算されていない	`zero_grad -> forward -> loss -> backward -> step` の順にする
lossがNaNになる	学習率、入力スケール、損失の定義域などが不安定	学習率を下げ、入力とlogの定義域を確認する
初期化後に学習済み重みを読み込む	読み込んだ重みを上書きしている	初期化してから重みを読み込む、または必要部分だけ初期化する
凍結したつもりの層が更新される	optimizerに更新対象として残っている	`requires_grad` とoptimizerのparamsを両方確認する
再開後に学習の挙動が変わる	optimizerやscheduler状態を保存していない	checkpointにまとめて保存する
KLDivLossの入力が確率	入力にはlog-probabilitiesが必要	`torch.log_softmax` を使う
CTCLossのshapeが合わない	`(N, T, C)` と `(T, N, C)` を混同している	log_probsのshapeを確認する

エラーが出たときは、関数名だけで悩まず、shape、dtype、値の範囲、device、勾配の有無を順番に見ます。損失関数とoptimizerの問題は、この基本確認でかなり切り分けられます。

参考リンク

torch.nn.init: https://docs.pytorch.org/docs/stable/nn.init.html
Loss Functions: https://docs.pytorch.org/docs/stable/nn.html#loss-functions
nn.CrossEntropyLoss: https://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html
nn.BCEWithLogitsLoss: https://docs.pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html
nn.KLDivLoss: https://docs.pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html
nn.CTCLoss: https://docs.pytorch.org/docs/stable/generated/torch.nn.CTCLoss.html
torch.optim: https://docs.pytorch.org/docs/stable/optim.html
torch.optim.SGD: https://docs.pytorch.org/docs/stable/generated/torch.optim.SGD.html
torch.optim.Adam: https://docs.pytorch.org/docs/stable/generated/torch.optim.Adam.html
torch.optim.AdamW: https://docs.pytorch.org/docs/stable/generated/torch.optim.AdamW.html
clip_grad_norm_: https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html

おわりに

今回は、モデルを訓練するために必要な、重み初期化、損失関数、最適化手法を整理しました。

重み初期化は、訓練の出発点を整えるためのものです。Xavier初期化は線形やTanh系の考え方と相性がよく、Kaiming初期化はReLU系のネットワークでよく使われます。model.apply() を使うと、層の種類ごとに初期化を分けられます。

損失関数では、モデル出力とtargetの対応が最重要です。単一ラベル多クラス分類では CrossEntropyLoss にlogitとクラス番号を渡します。二分類や多ラベル分類では、BCEWithLogitsLoss にlogitとfloatの0/1 targetを渡します。回帰や分布近似では、それぞれ別の損失を使います。

optimizerは、loss.backward() で計算された勾配を読み、パラメータを更新します。zero_grad -> forward -> loss -> backward -> step の順番を体で覚えると、訓練ループの見通しがよくなります。さらに、学習率、weight decay、パラメータグループ、勾配クリッピング、チェックポイントを組み合わせることで、実践的な訓練に近づきます。

次回は、ここまでの訓練ループをさらに実践に近づけるために、学習率スケジューリング、評価指標、ログ記録、過学習の見方、モデル保存と推論の流れを扱います。

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up

PyTorch入門 第5回：重み初期化、損失関数、最適化手法で学習を安定させる