Colab TPUでTensorFlowの重みを保存する方法

Google ColabのTPUランタイムでは、Google Cloud Strage(GCS)を使わない限り、TensorFlow公式の重み保存方法が使えません。そこで、TensorFlow公式の重み保存方法に頼らずに、重み保存と復元を行ってみます。

この記事はすべてGoogle ColabのTPUランタイムでTensorFlowを動かすことを前提としています。TPUに対応したTensorFlowの書き方はこちらの記事が詳しいです。
記事内のコードと実行結果はgithubでもColab Notebook形式で公開しています。

環境

• Google Colab (TPUランタイム)
• Python 3.6.8
• TensorFlow 2.0.0

TensorFlowの重み保存方法

本題に入る前にTensorFlowの重み保存方法をおさらいしておきます。

Checkpoint

TensorFlowで重み保存を行う1つ目の方法はtf.train.Checkpointです。こちらはTF1.xにおけるtf.train.saverを使った保存方法のTF2.0対応版です。
使い方は、モデル定義とオプティマイザー定義の後、学習の前でtf.train.Checkpointインスタンスを作成して、任意の場所でsaveメソッドを呼び出すだけ、と簡単です。さらにtf.train.CheckpointManagerでwrapしてあげることで、epoch毎の結果を簡単に保存・管理することができます。
詳しい使い方は公式ガイドを参照してください。

import tensorflow as tf

# モデル定義やら
# ...
# ...
model = tf.keras.Model(inputs=inputs, outputs=outputs)
optimizer = tf.optimizers.Adam(1.0e-4)

# checkpointオブジェクトを作成
ckpt = tf.train.Checkpoint(model=model, optimizer=optimizer)

# 学習
# ...
# ...

# 重み保存
ckpt.save(file_prefix='/tmp/save_test')

keras.Model.save

tf.train.Checkpoint以外の方法として、TF1.x時代からのtf.keras.Modelのsave_weightsメソッドを使ったやり方もあります。
使い方は、任意の場所でtf.keras.Modelインスタンスのsave_weightsメソッドを呼び出すだけです。こちらもとても簡単ですね。tf.keras.Modelインスタンスを使いますが、各レイヤーに紐づけられたオプティマイザーの各重み（Adamの1次モーメントや2次モーメント）も保存されるようです。

import tensorflow as tf

# モデル定義やら
# ...
# ...
model = tf.keras.Model(inputs=inputs, outputs=outputs)
optimizer = tf.optimizers.Adam(1.0e-4)

# 学習
# ...
# ...

# 重み保存
model.save_weights(file_path='/tmp/save_test')

TPUで動かしてみると…

これらの方法を使ってTPUで重み保存を実行してみます。
おそらく以下のようなErrorを吐いて保存できないはずです。

UnimplementedError: File system scheme '[local]' not implemented (file: '/content/save_test-1_temp_...')
Additional GRPC error information:
{"created":"...","description":"Error received from peer","file":"external/grpc/src/core/lib/surface/call.cc","file_line":1039,"grpc_message":"File system scheme '[local]' not implemented (file: '/content/save_test-1_temp_...')","grpc_status":12} [Op:SaveV2]

このErrorは、tf.data APIを使ってデータセットをtfrecordファイルから読み込もうとした場合に表示されるものと同じで、GCSを使う以外に直接の回避策はありません。(TPUのトラブルシューティング参照)

重みをNumpy配列で取り出しPickle漬けで保存

以上より、公式のAPIを使った保存方法はColab TPUでは使えません。そこで、自分で重みを取り出して保存する、というやり方が必要になります。

モデルの重みを取り出す

モデルをtf.keras.Modelを使って作成している場合、weights属性を参照することで、その時点でモデルに含まれる全ての重みを取り出すことができます。取り出した重みはtf.Variableオブジェクトですが、numpyメソッドでnumpy配列に変換できます。この方法はSequential APIで作成したモデル、Functional APIで作成したモデル、tf.keras.Modelを継承したクラスで作成したモデル、どれでもOKです。
保存する時は、後述の復元の時のために、変数名が分かるようなdictにしておくと便利です。

def get_model_weights_as_numpy(model):
    weights = {}
    for v in model.weights:
        # model.weightsで各Layerの重みを取り出し
        # 各variableはnumpyメソッドでnumpy配列に変換できる
        weights[v.name] = v.numpy()
    return {'model_name': model.name, 'weights': weights}

Optimizerの重みを取り出す

モデルの重みは保存できるような形で取り出せましたが、この中にオプティマイザーの重み（Adamの1次モーメントや2次モーメント）は含まれていません。これらはSlotと呼ばれ、別途取り出す必要があります。
tf.optimizers.Optimizerに設定されているSlotの名前はget_slot_namesメソッドで取り出し、モデルの各重みに紐づけられたSlotはget_slotメソッドで取り出すことができます。
モデルの重みと同様、復元する時のために、変数名とオプティマイザーの名前が分かるようにdictにしておくと便利です。

def get_optimizer_weights_as_numpy(optimizer, model):
    weights = {}
    slot_names = optimizer.get_slot_names()
    for v in model.weights:
        # model.weightsで各Layerの重みを取り出し
        weights[v.name] = {}
        for slot in slot_names:
            # 各Slotに対し、optimizerのget_slotで値を取り出す
            weights[v.name][slot] = optimizer.get_slot(v, slot).numpy()
    return {'optimizer_name': optimizer._name, 'weights': weights}

Pickle漬け

取り出した重みはPickle漬けにしてしまいましょう。モデルとオプティマイザーを別々に保存してもいいですが、管理が面倒なのでまとめて保存します。
今回はPickle漬けにして保存しますが、取り出した重みはNumpy配列となるので、他の方法でも保存が可能です。

import pickle

def save_weights_as_pickle(file_prefix, optimizer, model):
    model_weights = get_model_weights_as_numpy(model)
    optimizer_weights = get_optimizer_weights_as_numpy(optimizer, model)
    all_weights = {'model': model_weights, 'optimizer': optimizer_weights}

    with open(file_prefix + '.pkl', 'wb') as f:
        pickle.dump(all_weights, f)

Pickleを読み込んで重み復元

重みの保存はできましたが、保存したファイルから復元できないと保存した意味がありません。
基本方針は、重みが作成された後かつ学習を行う前に、各重みに対してassignメソッドで学習済み重みを割り当てる形です。
まずPickleから保存済みの重みを読み込んでおきましょう。

import pickle

def load_weights_from_pickle(file_prefix, optimizer, model):
    with open(file_prefix + '.pkl', 'rb') as f:
        weights = pickle.load(f)

    # modelの重みを復元（後述）
    set_model_weights_from_numpy(weights['model'], model)

    # optimizerの重みを復元（後述）
    set_optimizer_weights_from_numpy(weights['optimizer'], optimizer, model)

モデルの重み作成タイミングの問題

TF1.xと異なり、TF2.0では重みの作成タイミングが場合によって異なります。具体的には、各レイヤーの重みはレイヤーをインスタンス化した時点では作成されず、最初にcallされた時にbuildメソッドが呼び出されて重みが作成されます。（自作のレイヤーに関してはこの限りではありません）

Sequential APIまたはFunctional APIの場合

tf.keras.ModelをSequential APIまたはFunctional APIで作成した場合、各レイヤーはインスタンス化されて一度callされているので、モデルの作成時点で重みは作成されています。

import tensorflow as tf

input_shape = (28, 28, 1) # 入力のshape. 最初の次元（バッチサイズ）は除く.

# ネットワークの定義
# 入力層
x = tf.keras.layers.Input(input_shape)
# 畳み込み層1
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(x)
h = tf.keras.layers.ReLU()(h)
# 畳み込み層2
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(h)
h = tf.keras.layers.ReLU()(h)
# 畳み込み層3
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(h)
h = tf.keras.layers.ReLU()(h)
# 線形層
h = tf.keras.layers.Flatten()(h)
y = tf.keras.layers.Dense(10)(h)

# モデルの作成
model = tf.keras.Model(x, y)

# modelが保持している重みの名前を表示
print([v.name for v in model.weights])

['conv2d/kernel:0', 'conv2d/bias:0', 'conv2d_1/kernel:0', 'conv2d_1/bias:0', 'conv2d_2/kernel:0', 'conv2d_2/bias:0', 'dense/kernel:0', 'dense/bias:0']

この場合、モデルを作成しさえすれば、各重みに対してassignメソッドで学習済みの重みを割り当てることができるので、特に問題はありません。

Custom Modelの場合

tf.keras.Modelを継承したclassによるCustom Modelの場合、モデルの作成時点ではレイヤーがインスタンス化されるのみでcallされないので、重みは作成されません。

import tensorflow as tf

input_shape = (28, 28, 1) # 入力のshape. 最初の次元（バッチサイズ）は除く.

class CustomModel(tf.keras.Model):
    def __init__(self, *args, **kwargs):
        super(CustomModel, self).__init__(*args, **kwargs)

        # 畳み込み層1
        self.conv0 = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')
        self.act0 = tf.keras.layers.ReLU()
        # 畳み込み層2
        self.conv1 = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')
        self.act1 = tf.keras.layers.ReLU()
        # 畳み込み層3
        self.conv2 = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')
        self.act2 = tf.keras.layers.ReLU()
        # 線形層
        self.flatten = tf.keras.layers.Flatten()
        self.dense = tf.keras.layers.Dense(10)

    def call(self, inputs):
        h = self.conv0(inputs)
        h = self.act0(h)
        h = self.conv1(h)
        h = self.act1(h)
        h = self.conv2(h)
        h = self.act2(h)
        h = self.flatten(h)
        outputs = self.dense(h)
        return outputs

# モデルの作成
model = CustomModel()

# モデル作成の時点で保持している重み
print([v.name for v in model.weights])

x = tf.zeros((1, *input_shape), tf.float32)
y = model(x)

# 一度callした後に保持している重み
print([v.name for v in model.weights])

[]
['custom_model/conv2d/kernel:0', 'custom_model/conv2d/bias:0', 'custom_model/conv2d_1/kernel:0', 'custom_model/conv2d_1/bias:0', 'custom_model/conv2d_2/kernel:0', 'custom_model/conv2d_2/bias:0', 'custom_model/dense/kernel:0', 'custom_model/dense/bias:0']

そのため、モデルを作成した時点では重みが存在せずassignメソッドが使えないので、学習済みの重みを割り当てることができません。よってCustom Modelの場合は、学習前に一度callして重みを作成してから、学習済み重みを割り当てる必要があります。
学習前にテストデータを与えて推論を行う形でもいいですが、今回はCustom Modelを作成した時点で重みが作成されるように、__init__メソッド内でcallメソッドを呼び出してみます。

import tensorflow as tf

input_shape = (28, 28, 1) # 入力のshape. 最初の次元（バッチサイズ）は除く.

class CustomModel(tf.keras.Model):
    def __init__(self, input_shape, *args, **kwargs):
        super(CustomModel, self).__init__(*args, **kwargs)

        # 畳み込み層1
        self.conv0 = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')
        self.act0 = tf.keras.layers.ReLU()
        # 畳み込み層2
        self.conv1 = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')
        self.act1 = tf.keras.layers.ReLU()
        # 畳み込み層3
        self.conv2 = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')
        self.act2 = tf.keras.layers.ReLU()
        # 線形層
        self.flatten = tf.keras.layers.Flatten()
        self.dense = tf.keras.layers.Dense(10)

        # callして重みを作成する
        # dummy_inputsは1次元目(batch size)がNoneで2次元目以降がinput_shapeのTensor
        dummy_inputs = tf.keras.layers.Input(input_shape)
        _ = self.call(dummy_inputs)

    def call(self, inputs):
        h = self.conv0(inputs)
        h = self.act0(h)
        h = self.conv1(h)
        h = self.act1(h)
        h = self.conv2(h)
        h = self.act2(h)
        h = self.flatten(h)
        outputs = self.dense(h)
        return outputs

# モデルの作成
model = CustomModel(input_shape)

# modelが保持している重みの名前を表示
print([v.name for v in model.weights])

['conv2d/kernel:0', 'conv2d/bias:0', 'conv2d_1/kernel:0', 'conv2d_1/bias:0', 'conv2d_2/kernel:0', 'conv2d_2/bias:0', 'dense/kernel:0', 'dense/bias:0']

これで無事に学習前に重みを作成することができました。

モデルの重みを復元

あとは各重みに対してassignメソッドを使って学習済みの重みを割り当てるだけです。
重みの名前をkeyとしたdictで保存していたので、モデルの重みの名前と一致するものを復元します。

def set_model_weights_from_numpy(weights, model):
    for v in model.weights:
        if v.name in weights.keys():
            v.assign(weights[v.name])
        else:
            print('Not loaded weights: ' + v.name)

オプティマイザーの重み作成タイミングの問題

Custom Modelと同様に、オプティマイザーはインスタンス化した時でSlotが作成されません。インスタンス化の時点でモデルと紐づけされていないので当然といえば当然です。
オプティマイザーの場合は厄介なことに、通常Slotが作成されるのはBackpropが実行された時です。つまり、学習前にSlotを復元したいのに、学習を実行しないとSlotが作成されないわけです。

import tensorflow as tf

optimizer = tf.optimizers.Adam(1.0e-4)

# optimizerが保持している重みの名前を表示
print([v.name for v in optimizer.weights])

[]

そこで、tf.optimizers.Optimizerのadd_slotメソッドを使って、自動的にSlotが追加される前に手動でSlotを追加してしまいましょう。

import tensorflow as tf

input_shape = (28, 28, 1) # 入力のshape. 最初の次元（バッチサイズ）は除く.

# ネットワークの定義
# 入力層
x = tf.keras.layers.Input(input_shape)
# 畳み込み層1
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(x)
h = tf.keras.layers.ReLU()(h)
# 畳み込み層2
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(h)
h = tf.keras.layers.ReLU()(h)
# 畳み込み層3
h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(h)
h = tf.keras.layers.ReLU()(h)
# 線形層
h = tf.keras.layers.Flatten()(h)
y = tf.keras.layers.Dense(10)(h)

# モデルの作成
model = tf.keras.Model(x, y)

slot_names = ['m', 'v']
optimizer = tf.optimizers.Adam(1.0e-4)
with tf.name_scope(optimizer._name):
    for v in model.weights:
        for slot in slot_names:
            optimizer.add_slot(v, slot, initializer='zeros')

# optimizerが保持している重みの名前を表示
print([v.name for v in optimizer.weights])

['Adam/conv2d/kernel/m:0', 'Adam/conv2d/kernel/v:0', 'Adam/conv2d/bias/m:0', 'Adam/conv2d/bias/v:0', 'Adam/conv2d_1/kernel/m:0', 'Adam/conv2d_1/kernel/v:0', 'Adam/conv2d_1/bias/m:0', 'Adam/conv2d_1/bias/v:0', 'Adam/conv2d_2/kernel/m:0', 'Adam/conv2d_2/kernel/v:0', 'Adam/conv2d_2/bias/m:0', 'Adam/conv2d_2/bias/v:0', 'Adam/dense/kernel/m:0', 'Adam/dense/kernel/v:0', 'Adam/dense/bias/m:0', 'Adam/dense/bias/v:0']

これで学習前にオプティマイザーのSlotを作成することができました。

しかし、この方法でSlotを追加しても、実際には別のSlotが作成されたり、新しいSlotで上書きされていたら意味がありません。この方法で作成したSlotが無視されずに実際に使われるか、テストしてみましょう。
Adamオプティマイザーの1次モーメント$m$の初期値は通常ゼロですが、あえて1.0としてadd_slotメソッドでSlotを追加し、1回更新した後の値を確認してみます。$m$の更新式は以下の通りです。

m \leftarrow m \cdot \beta_1 + grad \cdot (1 - \beta_1)

$\beta_1$を0.9とした場合は以下のようになります。

m \leftarrow 0.9 \, m + 0.1 \, grad

import tensorflow as tf

input_shape = (10,)
learning_rate = 1.0e-4
beta_1 = 0.9
beta_2 = 0.99

inputs = tf.keras.layers.Input(input_shape)
outputs = tf.keras.layers.Dense(1, use_bias=False, kernel_initializer='ones')(inputs)
model = tf.keras.Model(inputs=inputs, outputs=outputs)
optimizer = tf.optimizers.Adam(learning_rate=learning_rate, beta_1=beta_1, beta_2=beta_2)

slot_names = ['m', 'v']
with tf.name_scope(optimizer._name):
    for v in model.weights:
        for slot in slot_names:
            # 本来はinitializerは'zeros'だが、あえて'ones'に設定
            # ここで追加されたslotが実際に使われていたら
            # 'zeros'の時と異なる値になるはず
            optimizer.add_slot(v, slot, initializer='ones')

x = tf.ones((1, *input_shape), tf.float32)
with tf.GradientTape() as tape:
    y = model(x)
    loss = 0.5 * tf.reduce_mean(y ** 2)
grad = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grad, model.trainable_variables))

# 勾配をチェック
print('=' * 20 + 'Gradient' + '=' * 20)
print(grad[0])
print()

print('=' * 20 + 'Actual 1st moment of Adam' + '=' * 20)
print(optimizer.get_slot(model.weights[0], 'm'))
print()

print('=' * 20 + 'estimated 1st moment of Adam' + '=' * 20)
print(1.0 * beta_1 + grad[0] * (1.0 - beta_1))
print()

====================Gradient====================
tf.Tensor(
[[10.]
 [10.]
 [10.]
 [10.]
 [10.]
 [10.]
 [10.]
 [10.]
 [10.]
 [10.]], shape=(10, 1), dtype=float32)

====================Actual 1st moment of Adam====================
<tf.Variable 'Adam/dense_12/kernel/m:0' shape=(10, 1) dtype=float32, numpy=
array([[1.9000002],
       [1.9000002],
       [1.9000002],
       [1.9000002],
       [1.9000002],
       [1.9000002],
       [1.9000002],
       [1.9000002],
       [1.9000002],
       [1.9000002]], dtype=float32)>

====================estimated 1st moment of Adam====================
tf.Tensor(
[[1.9]
 [1.9]
 [1.9]
 [1.9]
 [1.9]
 [1.9]
 [1.9]
 [1.9]
 [1.9]
 [1.9]], shape=(10, 1), dtype=float32)

もし、自動的に作成されたSlotが使われた場合、$m$の初期値はゼロなので、更新された値は $0.9 \times 0 + 0.1 \times 10.0 = 0.1$ となっているはずです。一方で、add_slotで追加したSlotが使われた場合、$m$の初期値を1.0とした時の値 $0.9 \times 1.0 + 0.1 \times 10.0 = 1.9$ となります。（3つ目のTensorの値）
結果を見ると、浮動小数点誤差っぽい値の違いはありますが、1.0で初期化されたSlotがbackpropで使われたことがわかります。
以上より、add_slotメソッドでSlotを追加することで重みが復元できそうです。

オプティマイザーの重みを復元

モデルの各重みに対して、保存済みの重みを初期値として指定してSlotを作成し、オプティマイザーの重みを復元します。
自動的に作成されたSlotはオプティマイザーの名前のscopeで作成されているので、add_slotメソッドで手動で作成する場合にも同じscopeを設定しておきます。

import tensorflow as tf

def set_optimizer_weights_from_numpy(weights, optimizer, model):
    # optimizerの名前でscopeする
    with tf.name_scope(weights['optimizer_name']):
        optimizer_weights = weights['weights']
        for v in model.weights:
            if v.name in optimizer_weights.keys():
                for slot in optimizer_weights[v.name].keys():
                    # 学習済みの重みを初期値としてslotを作成
                    initializer = tf.initializers.Constant(optimizer_weights[v.name][slot])
                    optimizer.add_slot(v, slot, initializer=initializer)
            else:
                print('Not loaded optimizer weights: ' + v.name)

CNNで実際にテストしてみる

CNNモデルを組んでMNISTを使ってテストしてみましょう。すべてGoogle ColabのTPUランタイムで実行します。使いやすくするためにTrainerクラスを定義します。

全体は長いので折りたたんでいます。

import os
import time
import pickle
import numpy as np
import tensorflow as tf

print(tf.__version__)

class Trainer(object):
    def __init__(self):
        self.batch_size = 256
        self.learning_rate = 1.0e-4
        self.input_shape = (28, 28, 1) # 入力のshape. 最初の次元（バッチサイズ）は除く.

        tpu_address = "grpc://" + os.environ["TPU_NAME"]
        cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu_address)
        tf.config.experimental_connect_to_cluster(cluster_resolver)
        tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
        self.tpu_strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)

        print('=' * 50)

        #=========================================================================
        # データセット
        #=========================================================================

        # データセットをロード
        # 今回はMNIST
        (X_train, Y_train), (X_test, Y_test) = tf.keras.datasets.mnist.load_data()
        X_train = X_train[..., np.newaxis].astype(np.float32)
        Y_train = Y_train.astype(np.int32)
        X_test = X_test[..., np.newaxis].astype(np.float32)
        Y_test = Y_test.astype(np.int32)
        self.N_train = X_train.shape[0]
        self.N_test = X_test.shape[0]

        with self.tpu_strategy.scope():
            # tf.data.Dataset APIを使う
            dataset_train = tf.data.Dataset.from_tensor_slices((X_train, Y_train))
            dataset_train = dataset_train.shuffle(buffer_size=self.N_train)
            dataset_train = dataset_train.batch(self.batch_size, drop_remainder=True)
            self.dataset_train = self.tpu_strategy.experimental_distribute_dataset(dataset_train)

            dataset_test = tf.data.Dataset.from_tensor_slices((X_test, Y_test))
            dataset_test = dataset_test.batch(self.batch_size, drop_remainder=True)
            self.dataset_test = self.tpu_strategy.experimental_distribute_dataset(dataset_test)

    def tpu_decorator(func):
        def wrapper(self, *args, **kwargs):
            if tf.distribute.in_cross_replica_context():
                outputs = func(self, *args, **kwargs)
            else:
                with self.tpu_strategy.scope():
                    outputs = func(self, *args, **kwargs)
            return outputs
        return wrapper

    @tpu_decorator
    def build_model(self):
        #=========================================================================
        # ネットワーク定義
        #=========================================================================

        # ネットワークの定義
        # 入力層
        x = tf.keras.layers.Input(self.input_shape)
        # 畳み込み層1
        h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(x)
        h = tf.keras.layers.ReLU()(h)
        # 畳み込み層2
        h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(h)
        h = tf.keras.layers.ReLU()(h)
        # 畳み込み層3
        h = tf.keras.layers.Conv2D(64, (3, 3), padding='SAME')(h)
        h = tf.keras.layers.ReLU()(h)
        # 線形層
        h = tf.keras.layers.Flatten()(h)
        y = tf.keras.layers.Dense(10)(h)

        # モデルの作成
        self.model = tf.keras.Model(x, y)
        self.optimizer = tf.optimizers.Adam(self.learning_rate)

    @tpu_decorator
    @tf.function
    def train_step(self, dist_inputs):
        def _train_step(inputs):
            images, labels = inputs

            # tf.GtadientTapeブロックで入力からロスまで計算
            with tf.GradientTape() as tape:
                logits = self.model(images)
                loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels, logits)
                loss = tf.reduce_sum(loss) / self.batch_size

            # gradientを計算
            grad = tape.gradient(loss, sources=self.model.trainable_variables)

            # optimizerで重みを更新
            self.optimizer.apply_gradients(zip(grad, self.model.trainable_variables))

            acc = tf.metrics.sparse_categorical_accuracy(labels, logits)
            acc = tf.reduce_sum(acc) / self.batch_size

            return loss, acc

        losses, accs = self.tpu_strategy.experimental_run_v2(_train_step, args=(dist_inputs,))
        losses = self.tpu_strategy.reduce(tf.distribute.ReduceOp.SUM, losses, axis=None)
        accs = self.tpu_strategy.reduce(tf.distribute.ReduceOp.SUM, accs, axis=None)
        return losses, accs

    @tpu_decorator
    @tf.function
    def eval_step(self, dist_inputs):
        def _eval_step(inputs):
            images, labels = inputs

            logits = self.model(images)
            loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels, logits)
            loss = tf.reduce_sum(loss) / self.batch_size
            acc = tf.metrics.sparse_categorical_accuracy(labels, logits)
            acc = tf.reduce_sum(acc) / self.batch_size
            return loss, acc

        losses, accs = self.tpu_strategy.experimental_run_v2(_eval_step, args=(dist_inputs,))
        losses = self.tpu_strategy.reduce(tf.distribute.ReduceOp.SUM, losses, axis=None)
        accs = self.tpu_strategy.reduce(tf.distribute.ReduceOp.SUM, accs, axis=None)
        return losses, accs

    @tpu_decorator
    def train(self, epochs):
        iterations = self.N_train // self.batch_size

        for epoch in range(epochs):
            time_start = time.time()
            train_loss = 0
            train_acc = 0
            dataset_iter = iter(self.dataset_train)

            for i in range(iterations):
                loss_tmp, acc_tmp = self.train_step(next(dataset_iter)) # 1step分の学習を実行
                train_loss += loss_tmp
                train_acc += acc_tmp

            # 平均ロスと平均精度
            epoch_loss = train_loss / iterations
            epoch_acc = 100 * train_acc / iterations

            # epochの結果を表示
            time_epoch = time.time() - time_start
            print('epoch: {:} loss: {:.4f} acc: {:.2f}% time: {:.2f}s'.format(
                epoch + 1, epoch_loss, epoch_acc, time_epoch))

    @tpu_decorator
    def eval(self):
        iterations = self.N_test // self.batch_size
        test_loss = 0
        test_acc = 0
        dataset_iter = iter(self.dataset_test)

        for i in range(iterations):
            loss_tmp, acc_tmp = self.eval_step(next(dataset_iter))
            test_loss += loss_tmp
            test_acc += acc_tmp
        test_loss /= iterations
        test_acc *= 100 / iterations
        print('test loss: {:.4f} acc: {:.2f}% '.format(test_loss, test_acc))

    @tpu_decorator
    def get_model_weights_as_numpy(self):
        weights = {}
        for v in self.model.weights:
            # model.weightsで各Layerの重みを取り出し
            # 各variableはnumpyメソッドでnumpy配列に変換できる
            weights[v.name] = v.numpy()
        return weights

    @tpu_decorator
    def get_optimizer_weights_as_numpy(self):
        weights = {}
        slot_names = self.optimizer.get_slot_names()
        for v in self.model.weights:
            # model.weightsで各Layerの重みを取り出し
            weights[v.name] = {}
            for slot in slot_names:
                # 各Slotに対し、optimizerのget_slotで値を取り出す
                weights[v.name][slot] = self.optimizer.get_slot(v, slot).numpy()
        return {'optimizer_name': self.optimizer._name, 'weights': weights}

    @tpu_decorator
    def save_weights_as_pickle(self, file_prefix):
        model_weights = self.get_model_weights_as_numpy()
        optimizer_weights = self.get_optimizer_weights_as_numpy()
        all_weights = {'model': model_weights, 'optimizer': optimizer_weights}

        with open(file_prefix + '.pkl', 'wb') as f:
            pickle.dump(all_weights, f)

    @tpu_decorator
    def set_model_weights_from_numpy(self, weights):
        for v in self.model.weights:
            if v.name in weights.keys():
                v.assign(weights[v.name])
            else:
                print('Not loaded weights: ' + v.name)

    @tpu_decorator
    def set_optimizer_weights_from_numpy(self, weights):
        # 必ずoptimizerの名前でscopeする
        with tf.name_scope(weights['optimizer_name']):
            optimizer_weights = weights['weights']
            for v in self.model.weights:
                if v.name in optimizer_weights.keys():
                    for slot in optimizer_weights[v.name].keys():
                        # 学習済みの重みを初期値としてslotを作成
                        initializer = tf.initializers.Constant(optimizer_weights[v.name][slot])
                        self.optimizer.add_slot(v, slot, initializer=initializer)
                else:
                    print('Not loaded optimizer weights: ' + v.name)

    @tpu_decorator
    def load_weights_from_pickle(self, file_prefix):
        with open(file_prefix + '.pkl', 'rb') as f:
            weights = pickle.load(f)

        self.set_model_weights_from_numpy(weights['model'])
        self.set_optimizer_weights_from_numpy(weights['optimizer'])

10エポック学習してテストデータにおけるロスと正解率を確認し、学習済みの重みを保存します。

trainer = Trainer()
trainer.build_model()
trainer.eval()
trainer.train(epochs=10)
trainer.eval()
trainer.save_weights_as_pickle('/content/save_test')

test loss: 6.4830 acc: 7.69% 
epoch: 1 loss: 0.2818 acc: 92.73% time: 12.25s
epoch: 2 loss: 0.0562 acc: 98.33% time: 7.11s
epoch: 3 loss: 0.0296 acc: 99.13% time: 7.21s
epoch: 4 loss: 0.0166 acc: 99.52% time: 6.97s
epoch: 5 loss: 0.0096 acc: 99.75% time: 6.83s
epoch: 6 loss: 0.0066 acc: 99.82% time: 7.28s
epoch: 7 loss: 0.0046 acc: 99.87% time: 7.01s
epoch: 8 loss: 0.0029 acc: 99.93% time: 7.30s
epoch: 9 loss: 0.0020 acc: 99.96% time: 7.08s
epoch: 10 loss: 0.0033 acc: 99.89% time: 6.98s
test loss: 0.0645 acc: 98.52%

Google Colabのランタイムを再起動し、もう一度Trainerクラスを定義しておきます。
学習済みの重みを復元してテストデータにおけるロスと正解率を見てみます。

trainer = Trainer()
trainer.build_model()
trainer.load_weights_from_pickle('/content/save_test')
trainer.eval()

test loss: 0.0645 acc: 98.52% 

学習直後の値と完全に一致しました。しっかりと重みを復元できていますね。

まとめ

Google ColabのTPUランタイムでは、TensorFlowの公式の重み保存方法は使えませんでした。
しかし、モデルの重みはnumpy配列で取得してPickle漬けにして保存、assignメソッドで復元することができました。
また、オプティマイザーの重みについてもget_slotメソッドで取得してPickle漬けにして保存、add_slotメソッドで復元することができました。

これでTPUランタイムで学習させた重みをローカルやGPUランタイム等で再利用することができますね。
記事内のコードと実行結果はgithubでもColab Notebook形式で公開しています。githubのリンクからGoogle Colabに飛ぶことができるので、気軽に試してみてください。