TensorFlow の自動微分 - GradientTape の使い方

ニューラルネットワークの訓練では、損失関数のパラメータに対する勾配を計算し、その勾配に基づいてパラメータを更新します。TensorFlow ではこの勾配計算を tf.GradientTape という仕組みで自動的に行えます。微分を手計算する必要がなく、複雑なモデルでも正確に勾配を求められるのが自動微分の強みです。

GradientTape の基本

tf.GradientTape はコンテキストマネージャとして使います。with ブロックの中で行われた計算を記録し、後から gradient メソッドで勾配を取り出せます。

import tensorflow as tf

x = tf.Variable(3.0)

with tf.GradientTape() as tape:
    y = x ** 2  # y = x^2

# dy/dx を計算
grad = tape.gradient(y, x)
print(grad)  # tf.Tensor(6.0, ...)

$y = x^{2}$ の導関数は $\frac{d y}{d x} = 2 x$ なので、 $x = 3$ のとき勾配は $6.0$ になります。GradientTape がこの微分を自動で計算してくれているわけです。

注意点として、GradientTape はデフォルトで tf.Variable のみを追跡します。tf.constant の勾配を求めたい場合は、明示的に tape.watch を呼ぶ必要があります。

x = tf.constant(3.0)

with tf.GradientTape() as tape:
    tape.watch(x)
    y = x ** 2

grad = tape.gradient(y, x)
print(grad)  # tf.Tensor(6.0, ...)

複数の変数に対する勾配

実際のモデルでは複数のパラメータに対する勾配を同時に計算します。gradient メソッドにリストを渡すことで、まとめて取得できます。

w = tf.Variable(2.0)
b = tf.Variable(1.0)

with tf.GradientTape() as tape:
    y = w * 3.0 + b  # y = 3w + b

grads = tape.gradient(y, [w, b])
print(grads[0])  # 3.0（dy/dw）
print(grads[1])  # 1.0（dy/db）

$y = 3 w + b$ の場合、 $\frac{\partial y}{\partial w} = 3$ 、 $\frac{\partial y}{\partial b} = 1$ となり、結果と一致しています。

テープの使い捨てと永続化

GradientTape はデフォルトでは1回の gradient 呼び出しでリソースが解放されます。同じ計算結果に対して複数回勾配を求めたい場合は persistent=True を指定します。

通常のテープ

gradient を1回呼ぶとテープが破棄される。メモリ効率が良い。通常の訓練ループではこちらで十分。

永続テープ

persistent=True で作成。複数回 gradient を呼べる。使い終わったら del tape で明示的に解放する必要がある。

x = tf.Variable(3.0)

with tf.GradientTape(persistent=True) as tape:
    y = x ** 2
    z = x ** 3

# 同じテープから複数の勾配を取得
dy_dx = tape.gradient(y, x)
dz_dx = tape.gradient(z, x)

print(dy_dx)  # 6.0
print(dz_dx)  # 27.0

del tape  # メモリ解放

高階微分

GradientTape を入れ子にすると、2階以上の微分も計算できます。たとえば $y = x^{3}$ の2階微分を求めるには、テープを二重にします。

x = tf.Variable(2.0)

with tf.GradientTape() as outer:
    with tf.GradientTape() as inner:
        y = x ** 3  # y = x^3

    # 1階微分: dy/dx = 3x^2
    dy_dx = inner.gradient(y, x)

# 2階微分: d^2y/dx^2 = 6x
d2y_dx2 = outer.gradient(dy_dx, x)

print(dy_dx)    # 12.0（3 * 2^2）
print(d2y_dx2)  # 12.0（6 * 2）

高階微分は物理シミュレーションや一部の高度な最適化手法で使われますが、通常のディープラーニングでは1階微分だけで十分なケースがほとんどです。

訓練ループでの実践的な使い方

GradientTape の本領が発揮されるのは、モデルの訓練ループです。以下は線形回帰を手動で訓練する例です。

import tensorflow as tf
import numpy as np

# ダミーデータ（y = 2x + 1 にノイズを加えたもの）
x_data = np.random.randn(100).astype(np.float32)
y_data = 2.0 * x_data + 1.0 + np.random.randn(100).astype(np.float32) * 0.1

# パラメータ
w = tf.Variable(0.0)
b = tf.Variable(0.0)
learning_rate = 0.1

for epoch in range(100):
    with tf.GradientTape() as tape:
        predictions = w * x_data + b
        loss = tf.reduce_mean((predictions - y_data) ** 2)

    grads = tape.gradient(loss, [w, b])
    w.assign_sub(learning_rate * grads[0])
    b.assign_sub(learning_rate * grads[1])

print(f"w = {w.numpy():.4f}")  # ≈ 2.0
print(f"b = {b.numpy():.4f}")  # ≈ 1.0

各エポックで GradientTape を使って損失の勾配を計算し、パラメータを手動で更新しています。実際のプロジェクトではオプティマイザの apply_gradients メソッドを使うのが一般的です。

オプティマイザとの組み合わせ

手動でのパラメータ更新は学習目的には良いですが、実用的にはオプティマイザを使います。Adam や SGD などのオプティマイザが、勾配に基づく更新を効率的に行ってくれます。

w = tf.Variable(0.0)
b = tf.Variable(0.0)
optimizer = tf.keras.optimizers.Adam(learning_rate=0.1)

for epoch in range(100):
    with tf.GradientTape() as tape:
        predictions = w * x_data + b
        loss = tf.reduce_mean((predictions - y_data) ** 2)

    grads = tape.gradient(loss, [w, b])
    optimizer.apply_gradients(zip(grads, [w, b]))

print(f"w = {w.numpy():.4f}")
print(f"b = {b.numpy():.4f}")

GradientTape で勾配を計算

optimizer.apply_gradients で更新

損失が収束するまで繰り返す

Keras の model.fit を使えばこの訓練ループ全体が自動化されますが、カスタムの損失関数や特殊な訓練手順が必要な場面では、GradientTape を直接使うことで柔軟に制御できます。

勾配が None になるケース

GradientTape を使っていて gradient が None を返すことがあります。よくある原因を把握しておくと、デバッグが楽になります。

変数が watch されていない

tf.constant を使っているのに tape.watch を忘れている場合。tf.Variable を使うか、明示的に watch を呼んでください。

計算がテープの外で行われた

with tf.GradientTape() ブロックの外で行った計算は記録されません。勾配を求めたい計算はすべてブロック内に含める必要があります。

整数型のテンソルを使っている

勾配は浮動小数点型にしか定義されません。tf.constant(3) ではなく tf.constant(3.0) のように、float 型を使ってください。

GradientTape は TensorFlow 2.x の訓練パイプラインの中核を成す仕組みです。Keras の model.fit だけでは対応できないカスタム訓練を実装する際に、この知識が役立ちます。