Python: TensorFlow2 の自動微分を試してみる

今回は、TensorFlow2 のプリミティブな API を使って、自動微分と勾配法で計算グラフを最適化する方法が気になったので試してみた。普段は Keras (tf.keras) を使ったミニバッチ学習をすることが多いけど、データのサイズが小さければバッチ学習で解く選択肢もあると思うので。

使った環境は次のとおり。

$ sw_vers
ProductName:    macOS
ProductVersion: 11.2
BuildVersion:   20D64
$ python -V     
Python 3.8.7
$ pip list | grep "tensorflow "
tensorflow             2.4.1

下準備

あらかじめ、TensorFlow をインストールしておく。

$ pip install tensorflow

サンプルコード

早速だけど以下にサンプルコードを示す。今回は単純すぎる例だけど x0^2 + x_1^2 を最小化してみた。 TensorFlow2 では GradientTape というコンテキストマネージャの中で Variable を操作すると、計算グラフが構築されて自動微分ができるようだ。勾配さえ計算できてしまえば、あとはそれをオプティマイザに放り込んでパラメータを更新するだけ。オプティマイザは特に何を使っても良いけどシンプルな例なので SGD を使った。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import annotations

import tensorflow as tf  # type: ignore


def objective(params: tf.Variable) -> tf.Tensor:
    """最小化したい目的関数"""
    # x_0^2 + x_1^2
    loss = params[0] ** 2 + params[1] ** 2
    return loss


@tf.function
def training_step(params: tf.Variable, optimizer: tf.keras.optimizers.Optimizer) -> None:
    # 自動微分する
    with tf.GradientTape(watch_accessed_variables=False) as t:
        # 勾配を計算するために追跡が必要なパラメータを登録する
        t.watch(params)  # watch_accessed_variables=True のときは不要
        # 計算グラフを構築する
        loss = objective(params)
    # 勾配を計算する
    grads = t.gradient(loss, params)
    # 勾配を使ってパラメータを更新する
    optimizer.apply_gradients([(grads, params)])
    # 現在のパラメータを出力する
    tf.print(params)


def main():
    # 定数で初期化した変数を用意する
    tensor = tf.constant([1., 4.], dtype=tf.float32)
    params = tf.Variable(tensor, trainable=True)

    # SGD で最適化する
    optimizer = tf.keras.optimizers.SGD(learning_rate=1e-1)

    # イテレーションを回す
    for _ in range(20):
        training_step(params, optimizer)

    # 結果を出力する
    print(objective(params))


if __name__ == '__main__':
    main()

上記を保存して実行してみる。 20 回イテレーションを回して、それぞれでパラメータの値を出力している。

$ python tfautograd.py 

...

[0.8 3.2]
[0.64 2.56]
[0.511999965 2.04799986]
[0.40959996 1.63839984]
[0.327679962 1.31071985]
[0.26214397 1.04857588]
[0.209715173 0.838860691]
[0.167772144 0.671088576]
[0.134217709 0.536870837]
[0.107374169 0.429496676]
[0.0858993381 0.343597353]
[0.068719469 0.274877876]
[0.0549755767 0.219902307]
[0.0439804606 0.175921842]
[0.0351843685 0.140737474]
[0.0281474944 0.112589978]
[0.0225179959 0.0900719836]
[0.0180143975 0.0720575899]
[0.0144115184 0.0576460734]
[0.0115292147 0.0461168587]
tf.Tensor(0.0022596875, shape=(), dtype=float32)