PyTorchにおける自動微分機能:計算グラフと勾配更新の実装詳細

オートグラッド(Autograd)の概要

PyTorch の学習プロセスを支える中核機能として、テンソル操作を追跡し勾配を自動的に計算する「Autograd」エンジンが存在します。この機構により、複雑なニューラルネットワークにおける微分計算を手動で行うことなく効率的に最適化が可能となります。

計算グラフの構築と追跡設定

テンソルの作成時に requires_grad=True を指定することで、PyTorch はその張量に対する演算履歴を記録し始めます。これにより実行時において動的な計算グラフが形成され、入力から出力への依存関係が管理されます。

import torch

# 勾配追跡を有効にした入力データを作成
input_tensor = torch.ones(2, dtype=torch.float32, requires_grad=True)
print(input_tensor.requires_grad)
# 結果: True

上記のように定義された変数は、後にバックプロパゲーションの対象となるリーフノードとして扱われます。

フォワードパスによるグラフ生成

勾配追跡中の変数に対して数学的演算を行うと、新たな結果テンソルが生成され、元の演算情報を保持したまま接続されます。

# 単純な線形変換を行う
weighted_output = input_tensor * 3.0 
scalar_loss = weighted_output.sum()

print(scalar_loss.grad_fn)
# 出力例: <SumBackward0 object at 0x...>

ここで生じた grad_fn は、この値に至るまでの計算経路を示す関数オブジェクトです。

バックプロパゲーションによる勾配取得

損失関数(スカラー値)に対して .backward() メソッドを呼び出すことで、連鎖律に基づき各パラメータに対する偏微分値が計算され、該当するテンソルの .grad プロパティへ格納されます。

scalar_loss.backward()

# 入力テンソルに対する勾配を確認
print(input_tensor.grad)
# 結果: tensor([[3., 3.]])
# 計算: loss = sum(x * 3) = 3 * sum(x) => d(loss)/dx = 3

モデル訓練ループのパターン

実際のモデル学習では、予測値の算出、損失の計算、勾配の逆伝播、そしてパラメータ更新というサイクルを繰り返します。勾配はデフォルトで累積されるため、バッチ処理毎に初期化が必要です。

import torch.nn as nn
import torch.optim as optim

# サンプルデータ (y = 0.5x + 5 にフィットさせる)
x_data = torch.tensor([[1.0], [2.0], [3.0]], requires_grad=False)
y_target = torch.tensor([[5.5], [6.0], [6.5]], requires_grad=False)

# パラメータ初期化(学習対象)
weights = torch.randn(1, 1, requires_grad=True)
bias = torch.zeros(1, requires_grad=True)

# 最適化アルゴリズムの設定(SGD, lr=0.1)
optimizer = optim.SGD([weights, bias], lr=0.1)
criterion = nn.L1Loss()

for step in range(500):
    # 推論処理
    y_pred = x_data @ weights + bias
    
    # 損失計算
    loss = criterion(y_pred, y_target)
    
    # グラディエントの初期化
    optimizer.zero_grad()
    
    # 逆伝播計算
    loss.backward()
    
    # パラメータ更新
    optimizer.step()
    
    if step % 100 == 0:
        print(f'Step {step}, Loss: {loss.item():.4f}')

print(f'最終重み: {weights.item()}, 最終バイアス: {bias.item()}')

これにより、重みとバイアスが真の値である 0.5 および 5.0 に収束していく様子が確認できます。

性能最適化に関する注意点

  • メモリ効率: 評価フェーズや推論時など、勾配計算が不要な場面では torch.inference_mode() または torch.no_grad() コンテキストマネージャを使用すると、計算図の構築が回避されメモリ使用量が削減されます。
  • グラフの維持: 同一グラフ上で複数回のバックプロパゲーションが必要となる特殊なケースでは、backward(retain_graph=True) を指定することで計算図の破棄を防ぐことが可能です。
  • 高次微分: 二階微分などの高度な計算を行う際には create_graph=True オプションを設定して、勾配計算自体を微分可能なグラフの一部として扱い続けます。

タグ: PyTorch autograd backpropagation deep-learning optimization

7月3日 17:39 投稿