TensorFlow2の動作確認用のコードサンプル（非チュートリアル）

インストール直後にやるべき動作確認

TensorFlowをインストールしたあと、動作確認に使うコードのサンプルを紹介します。

• バージョン確認だけだと、学習させたときにエラーが出ないか、GPUを使っているかなどの確認ができない。
• チュートリアルだと長過ぎる(というかコピペできるように紹介してない)ことが多い。

ので、簡単な動作確認として、バージョンとGPU確認、最低限の学習をするコードを書いています。

import numpy as np
import tensorflow as tf

print(tf.__version__)

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print("Physical GPUs: {}, Logical GPUs: {}".format(len(gpus), len(logical_gpus)))
else:
    print("CPU only")

x = np.arange(-1, 1, 0.0001)
y = 0.8 * x + 0.2

model = tf.keras.Sequential([tf.keras.layers.Dense(1, activation=None)])
model.compile("sgd", "mse")
model.build(input_shape=(0,1))
model.summary()
model.fit(x, y, epochs=5)

print("ground truth: 0.8, 0.2")
print("estimated: ", model.variables[0][0,0].numpy(), model.variables[1][0].numpy())

バージョンやGPUの情報がTensorFlowのログに紛れて見にくいので、必要に応じて改行とかを入れてください。

dockerを使っていたりすると、GPUの設定を間違えやすいので、GPUが認識されているかはしっかり確認しましょう。
学習中にnvidia-smiでGPUが使われているかの確認もしておきましょう。
学習が始まるとGPU-UtilとMemory-Usageが増えて、学習が終わるとGPU-Utilだけ下がり、Memory-Usageはそのままになるはずです。

学習については、$y = ax + b$の$a(=0.8)$, $b=(0.2)$を線形回帰していて、estimatedとして近い値を求められていればOKです。