average_optimizers_callback.ipynb

Kernel: Python 3

Copyright 2020 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

モデルの平均化

TensorFlow.orgで 表示

Google Colab で実行

GitHub でソースを表示

ノートブックをダウンロード

概要

このノートブックでは、TensorFlow アドオンパッケージから移動平均オプティマイザとモデル平均チェックポイントを使用する方法を紹介します。

移動平均化

移動平均化の利点は、最新のバッチで激しい損失の変化や不規則なデータ表現が発生しにくいことです。ある時点までモデルのトレーニングがスムーズになり、より一般的なアイデアを提供します。

確率的平均化

確率的重み平均化は、より広いオプティマイザに収束します。これは幾何学的なアンサンブルに似ています。確率的重み平均化は、他のオプティマイザのラッパーとして使用し、内側のオプティマイザのトラジェクトリの異なる点からの結果を平均化することでモデルの性能を向上させる、シンプルな方法です。

モデル平均チェックポイント

callbacks.ModelCheckpointにはトレーニングの途中で移動平均の重みを保存するオプションがないため、モデル平均オプティマイザにはカスタムコールバックが必要でした。update_weightsパラメータを使用すると、ModelAverageCheckpointで以下が可能になります。

1. モデルに移動平均重みを割り当てて保存する。

2. 古い平均化されていない重みはそのままにして、保存されたモデルは平均化された重みを使用する。

セットアップ

!pip install -U tensorflow-addons

import tensorflow as tf
import tensorflow_addons as tfa

import numpy as np
import os

モデルを構築する

def create_model(opt):
    model = tf.keras.models.Sequential([
        tf.keras.layers.Flatten(),                         
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])

    model.compile(optimizer=opt,
                    loss='sparse_categorical_crossentropy',
                    metrics=['accuracy'])

    return model

データセットを準備する

#Load Fashion MNIST dataset
train, test = tf.keras.datasets.fashion_mnist.load_data()

images, labels = train
images = images/255.0
labels = labels.astype(np.int32)

fmnist_train_ds = tf.data.Dataset.from_tensor_slices((images, labels))
fmnist_train_ds = fmnist_train_ds.shuffle(5000).batch(32)

test_images, test_labels = test

ここでは、次の 3 つのオプティマイザを比較してみます。

• ラップされていない SGD

• 移動平均を適用した SGD

• 確率的重み平均を適用した SGD

同じモデルを使用してパフォーマンスを見てみましょう。

#Optimizers 
sgd = tf.keras.optimizers.SGD(0.01)
moving_avg_sgd = tfa.optimizers.MovingAverage(sgd)
stocastic_avg_sgd = tfa.optimizers.SWA(sgd)

MovingAverage と StochasticAverage オプティマイザはどちらも ModelAverageCheckpoint を使用します。

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

cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_dir,
                                                 save_weights_only=True,
                                                 verbose=1)
avg_callback = tfa.callbacks.AverageModelCheckpoint(filepath=checkpoint_dir, 
                                                    update_weights=True)

モデルをトレーニングする

Vanilla SGD オプティマイザ

#Build Model
model = create_model(sgd)

#Train the network
model.fit(fmnist_train_ds, epochs=5, callbacks=[cp_callback])

#Evalute results
model.load_weights(checkpoint_dir)
loss, accuracy = model.evaluate(test_images, test_labels, batch_size=32, verbose=2)
print("Loss :", loss)
print("Accuracy :", accuracy)

移動平均 SGD

#Build Model
model = create_model(moving_avg_sgd)

#Train the network
model.fit(fmnist_train_ds, epochs=5, callbacks=[avg_callback])

#Evalute results
model.load_weights(checkpoint_dir)
loss, accuracy = model.evaluate(test_images, test_labels, batch_size=32, verbose=2)
print("Loss :", loss)
print("Accuracy :", accuracy)

確率的重み平均 SGD

#Build Model
model = create_model(stocastic_avg_sgd)

#Train the network
model.fit(fmnist_train_ds, epochs=5, callbacks=[avg_callback])

#Evalute results
model.load_weights(checkpoint_dir)
loss, accuracy = model.evaluate(test_images, test_labels, batch_size=32, verbose=2)
print("Loss :", loss)
print("Accuracy :", accuracy)