PythonでKerasを使って特定のエポック数ごとにモデルの重みを保存する方法

TensorFlowはGoogleが提供する機械学習フレームワークです。オープンソースとして公開されており、Pythonと組み合わせてアルゴリズムやディープラーニングアプリケーションの実装など、幅広い用途に活用されています。研究用途から本番環境まで、さまざまな場面で採用されている実績のあるフレームワークです。

TensorFlowには最適化技術が組み込まれており、複雑な数値計算を高速に処理できる点も大きな特徴です。

「tensorflow」パッケージは、Windows環境では以下のコマンドでインストールできます。

pip install tensorflow

テンソル（Tensor）はTensorFlowにおける基本的なデータ構造です。データフローグラフと呼ばれる計算グラフ内のエッジ（辺）をつなぐ役割を担います。テンソルは多次元配列またはリストと考えると分かりやすいでしょう。

Kerasとは

Kerasは、ONEIROS（Open ended Neuro-Electronic Intelligent Robot Operating System）というプロジェクトの研究の一環として開発されました。Pythonで書かれたディープラーニングAPIであり、機械学習の問題を効率的に解決するための高水準インターフェースを提供します。

Kerasは高いスケーラビリティとクロスプラットフォーム対応力を持っています。TPUやGPUクラスタ上でも動作可能で、学習済みモデルはWebブラウザやモバイルデバイス向けにエクスポートすることもできます。

KerasはTensorFlowパッケージにあらかじめ含まれており、以下のコードで利用できます。

import tensorflow
from tensorflow import keras

以降のコードはGoogle Colaboratoryで実行しています。Google Colabを使うと、ブラウザ上でPythonコードを実行でき、事前設定は不要。GPUにも無料でアクセスできるため、手軽に試せるのが魅力です。ColaboratoryはJupyter Notebookをベースに構築されています。

チェックポイントコールバックの実装例

それでは、4エポックごとにモデルの重みを保存するコードを見てみましょう。

サンプルコード

checkpoint_path = "training_2/cp-{epoch:04d}.ckpt"
checkpoint_dir = os.path.dirname(checkpoint_path)

batch_size = 32
print("Callback being created to save the model's weight after every 4 epoch")
cp_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_path,
    verbose=1,
    save_weights_only=True,
    save_freq=4*batch_size)

print("A new model instance is created")
model = create_model()
print("The weights are saved using 'checkpoint_path'")
model.save_weights(checkpoint_path.format(epoch=0))

コード出典：https://www.tensorflow.org/tutorials/keras/save_and_load

実行結果

Callback being created to save the model's weight after every 4 epoch
A new model instance is created
The weight are saved using 'checkpoint_path'

コードの解説

• ModelCheckpointコールバック：チェックポイントに一意の名前を付けたり、保存頻度を調整したりなど、柔軟なオプションが多数用意されています。
• save_freqパラメータ：この例では「save_freq=4*batch_size」と指定することで、4エポックごとに重みが保存されるよう設定しています。
• save_weights_only=True：モデル全体ではなく重みのみを保存するため、ストレージ容量を節約できます。
• 新しいモデルの作成と初期保存：新しく作成したモデルインスタンスに対して、epoch=0の時点で重みを保存し、学習途中の状態から再開できるようにしています。

このようにModelCheckpointコールバックを活用すれば、長時間の学習中に予期せぬ中断が発生しても、直近の重みから学習を再開できます。特に大規模なデータセットでの学習において、非常に有用な機能といえるでしょう。