概要

KerasでGPUを使ったニューラルネットワークの学習において、再現性を確保するために苦労したのでその方法をまとめます。
最初に見つけた方法でひとまず目的は達成されたのですが、TensorFlow 2.8.0から実装された簡単な方法もありましたので、そちらも記載しました。

環境

執筆時点　2022/5/15
google colab
Python 3.7.13
TensorFlow ver. 2.8.0
GPU Tesla K80, T4

情報に過不足がありましたらご指摘よろしくお願いします。

結論

コードだけ知りたい方へ。下記をモデル作成前に追加することで再現が取れました。

TensorFlow 2.1以降

import os
import random
import numpy as np
import tensorflow as tf

os.environ['PYTHONHASHSEED'] = '0'
os.environ['TF_DETERMINISTIC_OPS'] = '1'

np.random.seed(1)
random.seed(1)
tf.random.set_seed(1)

TensorFlow 2.8.0以降

import tensorflow as tf

tf.keras.utils.set_random_seed(1)
tf.config.experimental.enable_op_determinism()

実験

下記のコードで確認しました。コード先頭で環境変数PYTHONHASHSEEDとop determinism有効化の設定、モデル作製の直前でシード値の設定を行っています。TensorFlow 2.8.0で使用可能なコードはコメントアウトしています。ノートブックはこちら。

Keras_Reproducibility.ipynb

import os
import random

import numpy as np
import tensorflow as tf


 # 再現性確保のための設定
 # TensorFlow 2.1以降
+ os.environ['PYTHONHASHSEED'] = '0'  # ハッシュベースの操作を再現可能にする
+ os.environ['TF_DETERMINISTIC_OPS'] = '1'  # op determinismを有効化

 # TensorFlow 2.8.0以降
 # tf.config.experimental.enable_op_determinism()  # op determinismを有効化


 # データの読み込み
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

 # 以下のモデルの学習が再現できるか確認
def fit_model(x_train, y_train):
    model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10),
    tf.keras.layers.Softmax()
    ])

    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

    model.compile(optimizer='adam',
                loss=loss_fn,
                metrics=['accuracy'])

    history = model.fit(x_train, y_train, epochs=3, validation_split=0.2)

    return history


 # モデルの学習を2回行って比較
 # ループ内のモデル作製前にシード値を設定
histories = []

for i in range(2):
    # シード値の設定
    # TensorFlow 2.1以降
+    np.random.seed(1)
+    random.seed(1)
+    tf.random.set_seed(1)
    
    # TensorFlow 2.8.0以降
    # tf.keras.utils.set_random_seed(1)


    print(f'Model {i}')
    histories.append(fit_model(x_train, y_train))


 # 結果の確認
print('')
print('Result: {}'.format(histories[0].history == histories[1].history))

TensorFlow ver. 2.1以降のコード詳細

環境変数 PYTHONHASHSEED の指定

Python 3.2.3以降でハッシュベースの操作を再現可能にするために必要な設定とのことです。'0'に指定することで、ハッシュのランダム化が無効化されます。コードの先頭で一度設定すれば大丈夫でした。
上記の実験ではこの環境変数をいじらなくても学習結果の再現ができましたが、ドキュメントの指示に従いました。

os.environ['PYTHONHASHSEED'] = '0'

Keras documentation

First, you need to set the PYTHONHASHSEED environment variable to 0 before the program starts (not within the program itself). This is necessary in Python 3.2.3 onwards to have reproducible behavior for certain hash-based operations (e.g., the item order in a set or a dict, see Python's documentation or issue #2280 for further details).

Python documentation

If this variable is not set or set to random, a random value is used to seed the hashes of str, bytes and datetime objects.
If PYTHONHASHSEED is set to an integer value, it is used as a fixed seed for generating the hash() of the types covered by the hash randomization.
Its purpose is to allow repeatable hashing, such as for selftests for the interpreter itself, or to allow a cluster of python processes to share hash values.
The integer must be a decimal number in the range [0,4294967295]. Specifying the value 0 will disable hash randomization.

op determinism の有効化

TensorFlow 2.1以降でGPU-deterministic op functionalityを使用する設定です。下記によりGPUでの計算順序が固定されます。コードの先頭で一度設定すれば大丈夫でした。

os.environ['TF_DETERMINISTIC_OPS'] = '1'

Kerasのドキュメントによれば、通常のGPUの並列計算では実行順序が保証されず、浮動小数点数の精度の問題から計算順序が変わることで微小な誤差が生じるようです。（CPUを使えって言いますがそれは・・）

Moreover, when running on a GPU, some operations have non-deterministic outputs, in particular tf.reduce_sum(). This is due to the fact that GPUs run many operations in parallel, so the order of execution is not always guaranteed. Due to the limited precision of floats, even adding several numbers together may give slightly different results depending on the order in which you add them. You can try to avoid the non-deterministic operations, but some may be created automatically by TensorFlow to compute the gradients, so it is much simpler to just run the code on the CPU.

シード値の設定

上から順にPythonで生成するコアな乱数、Numpyで生成する乱数、TensorFlowで生成する乱数のシード値の設定です。上記の実験では上2つは設定せずとも再現できましたが、ドキュメントに従って実行しています。op determinismを有効化したときは、tf.random.set_seed(1)を実行しないとエラーが出ます。

random.seed(1)
np.random.seed(1)
tf.random.set_seed(1)

TensorFlow ver. 2.8.0以降のコード詳細

op determinism の有効化

以前のバージョンでは環境変数をいじっていましたが、関数一つで有効化できるようになっていました。

tf.config.experimental.enable_op_determinism()

シード値の設定

以前のバージョンではPython、Numpy、TensorFlowのシード値を別々の関数で設定していましたが、その機能を一つにまとめた関数です。

tf.config.experimental.enable_op_determinism()

補足

2.8.0では特に記載がなかったためPYTHONHASHSEEDの設定はしていませんが、モデルの再現は取れました。再現性に問題が生じたらPYTHONHASHSEED設定を試してみるのが良いかもしれません。

まだdeterministicな実装が行われていない操作もあるため、その場合はUnimplementedErrorが発生するようです。またバグのためにUnimplementedErrorを返さないことがあれば、報告して欲しいとのことでした。

TensorFlow documentation

Certain ops will raise an UnimplementedError because they do not yet have a deterministic implementation. Additionally, due to bugs, some ops might be nondeterministic and not raise an UnimplementedError. If you encounter such ops, please file an issue.

参考