PyTorch の自動微分 - autograd の仕組み

PyTorch が機械学習フレームワークとして広く使われている理由のひとつが、自動微分エンジン autograd です。ニューラルネットワークの学習には勾配（微分値）の計算が不可欠ですが、autograd はこれを自動的に処理してくれます。

autograd とは

ニューラルネットワークの学習では、損失関数の値を小さくするためにパラメータを少しずつ調整します。この調整方向を決めるのが勾配で、勾配を求める操作が微分です。

autograd は、テンソルに対する演算の履歴を内部的に記録し、.backward() を呼ぶだけで勾配を自動計算してくれる仕組みです。手動で微分の数式を書く必要はありません。

import torch

x = torch.tensor(3.0, requires_grad=True)
y = x ** 2 + 2 * x + 1

y.backward()
print(x.grad)  # tensor(8.)

この例では $y=x^2+2x+1$ を定義し、backward() で $x$ に関する微分を計算しています。$\frac{dy}{dx}=2x+2$ なので、$x=3$ のとき勾配は $8$ になります。

requires_grad の役割

requires_grad=True を指定したテンソルは、そこから派生するすべての演算が追跡されます。逆に、入力データなど勾配が不要なテンソルにはこのフラグを立てる必要はありません。

計算グラフの仕組み

autograd の内部では、演算のたびに計算グラフ（有向非巡回グラフ）が構築されます。各ノードは演算を表し、backward() を呼ぶとこのグラフを末端から入力方向へたどりながら連鎖律（チェインルール）で勾配を伝播させます。

a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(3.0, requires_grad=True)

c = a * b
d = c + a
d.backward()

print(a.grad)  # tensor(4.)  dc/da + dd/da = b + 1 = 4
print(b.grad)  # tensor(2.)  dc/db = a = 2

変数 d から a に至るまでの経路が複数ある場合、各経路の勾配が合算されます。これが計算グラフの強みで、複雑なネットワークでも正確に勾配を求められます。

勾配の蓄積と初期化

PyTorch では backward() を呼ぶたびに勾配が加算されます。これは一部のアルゴリズムでは便利ですが、通常の学習ループでは意図しない動作を引き起こすため、毎回 zero_() で初期化する必要があります。

x = torch.tensor(1.0, requires_grad=True)

for i in range(3):
    y = x * 2
    y.backward()
    print(f"ループ {i}: grad = {x.grad}")
    x.grad.zero_()  # これがないと勾配が蓄積される

optimizer.zero_grad() という形で使うのが一般的で、学習ループの冒頭で呼び出すのが定番のパターンです。

torch.no_grad() で推論を高速化

モデルの推論時や評価時には勾配計算が不要です。torch.no_grad() コンテキスト内では計算グラフが構築されないため、メモリ消費を抑えつつ処理速度を向上できます。

model = torch.nn.Linear(10, 1)
x = torch.randn(5, 10)

# 学習時
output = model(x)  # 計算グラフが構築される

# 推論時
with torch.no_grad():
    output = model(x)  # 計算グラフは構築されない

学習と推論を切り替える場面は頻繁にあるため、torch.no_grad() は PyTorch を使ううえで欠かせないイディオムです。

実践的な学習ループでの autograd

ここまでの要素をまとめると、典型的な学習ループは次のような流れになります。

実際のコードで確認してみましょう。

import torch
import torch.nn as nn

model = nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()

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

for epoch in range(100):
    optimizer.zero_grad()
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()
    optimizer.step()

print(f"weight: {model.weight.item():.3f}, bias: {model.bias.item():.3f}")

この例は $y=2x$ の関係を学習する最小限のコードです。autograd が勾配計算を担い、オプティマイザがパラメータ更新を行うという役割分担が PyTorch の設計思想の根幹になっています。