TensorFlow2.0＋TPUで複数のモデルを交互に訓練するテスト

Posted On 2019-10-17

614{icon} {views}

書籍情報

技術書典8の新刊『モザイク除去から学ぶ　最先端のディープラーニング』（A4・195ページ）好評通販中です！　機械学習の入門からGANの最先端までを書いたおすすめの本となっています！　Boothで試し読みできます。情報まとめ・質問用GitHub

GANの利用を想定します。以前TPUだと複数のモデルを同時or交互に訓練するというのは厳しかったのですが、これがTF2.0で変わったのか確かめます。

環境：TensorFlow2.0.0、Colab TPU

コード

おまじない

import tensorflow as tf
import os
# tpu用
# 詳細 https://www.tensorflow.org/guide/distributed_training#tpustrategy
tpu_grpc_url = "grpc://" + os.environ["COLAB_TPU_ADDR"]
tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
tf.config.experimental_connect_to_cluster(tpu_cluster_resolver) # TF2.0の場合、ここを追加
tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver) # TF2.0の場合、今後experimentialが取れる可能性がある    
strategy = tf.distribute.experimental.TPUStrategy(tpu_cluster_resolver)  # ここも同様

KerasのAPIを使うのではなくカスタムの訓練ループを書きます。

import tensorflow as tf
import tensorflow.keras.layers as layers

def model_a():
    inputs = layers.Input((1,))
    x = layers.Dense(2, use_bias=False)(inputs)
    return tf.keras.models.Model(inputs, x)

def model_b():
    inputs = layers.Input((2,))
    x = layers.Dense(1, use_bias=False)(inputs)
    return tf.keras.models.Model(inputs, x)

def main():
    with strategy.scope():
        ma = model_a()
        mb = model_b()
        loss_func = tf.keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.NONE)
        optima = tf.keras.optimizers.Adam()
        optimb = tf.keras.optimizers.Adam()

        print("Before update")
        print(ma.trainable_weights)
        print(mb.trainable_weights)

        def train_a():
            # update A
            with tf.GradientTape() as tape:
                x = tf.random.normal(shape=(2,1))
                y_pred = mb(ma(x))
                y = tf.ones(shape=(2, 1))
                loss = loss_func(y, y_pred)
            grad = tape.gradient(loss, ma.trainable_weights)
            optima.apply_gradients(zip(grad, ma.trainable_weights))

        def train_b():
            # update B
            with tf.GradientTape() as tape:
                x = tf.random.normal(shape=(2,1))
                y_pred = mb(ma(x))
                y = tf.ones(shape=(2, 1))
                loss = loss_func(y, y_pred)
            grad = tape.gradient(loss, mb.trainable_weights)
            optimb.apply_gradients(zip(grad, mb.trainable_weights))

        @tf.function
        def distributed_train_a():
            return strategy.experimental_run_v2(train_a)

        @tf.function
        def distributed_train_b():
            return strategy.experimental_run_v2(train_b)

        distributed_train_a()

        print("Update only A ")
        print(ma.trainable_weights)
        print(mb.trainable_weights)

        distributed_train_b()

        print("Update only B")
        print(ma.trainable_weights)
        print(mb.trainable_weights)

if __name__ == "__main__":
    main()

「ModelB(ModelA(input))」のような出力をして、あとは適当に損失関数を入れ、「Aのみ訓練→Bのみ訓練」がうまく行っているか確かめます。

結果

Before update
[TPUMirroredVariable:{
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense/kernel:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.31612158, -0.6931683 ]], dtype=float32)>,
  1 /job:worker/replica:0/task:0/device:TPU:1: <tf.Variable 'dense/kernel/replica_1:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.31612158, -0.6931683 ]], dtype=float32)>,
  2 /job:worker/replica:0/task:0/device:TPU:2: <tf.Variable 'dense/kernel/replica_2:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.31612158, -0.6931683 ]], dtype=float32)>,
  3 /job:worker/replica:0/task:0/device:TPU:3: <tf.Variable 'dense/kernel/replica_3:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.31612158, -0.6931683 ]], dtype=float32)>,
  4 /job:worker/replica:0/task:0/device:TPU:4: <tf.Variable 'dense/kernel/replica_4:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.31612158, -0.6931683 ]], dtype=float32)>,
  5 /job:worker/replica:0/task:0/device:TPU:5: <tf.Variable 'dense/kernel/replica_5:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.31612158, -0.6931683 ]], dtype=float32)>,
  6 /job:worker/replica:0/task:0/device:TPU:6: <tf.Variable 'dense/kernel/replica_6:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.31612158, -0.6931683 ]], dtype=float32)>,
  7 /job:worker/replica:0/task:0/device:TPU:7: <tf.Variable 'dense/kernel/replica_7:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.31612158, -0.6931683 ]], dtype=float32)>
}]
[TPUMirroredVariable:{
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense_1/kernel:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  1 /job:worker/replica:0/task:0/device:TPU:1: <tf.Variable 'dense_1/kernel/replica_1:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  2 /job:worker/replica:0/task:0/device:TPU:2: <tf.Variable 'dense_1/kernel/replica_2:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  3 /job:worker/replica:0/task:0/device:TPU:3: <tf.Variable 'dense_1/kernel/replica_3:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  4 /job:worker/replica:0/task:0/device:TPU:4: <tf.Variable 'dense_1/kernel/replica_4:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  5 /job:worker/replica:0/task:0/device:TPU:5: <tf.Variable 'dense_1/kernel/replica_5:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  6 /job:worker/replica:0/task:0/device:TPU:6: <tf.Variable 'dense_1/kernel/replica_6:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  7 /job:worker/replica:0/task:0/device:TPU:7: <tf.Variable 'dense_1/kernel/replica_7:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>
}]
Update only A 
[TPUMirroredVariable:{
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense/kernel:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  1 /job:worker/replica:0/task:0/device:TPU:1: <tf.Variable 'dense/kernel/replica_1:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  2 /job:worker/replica:0/task:0/device:TPU:2: <tf.Variable 'dense/kernel/replica_2:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  3 /job:worker/replica:0/task:0/device:TPU:3: <tf.Variable 'dense/kernel/replica_3:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  4 /job:worker/replica:0/task:0/device:TPU:4: <tf.Variable 'dense/kernel/replica_4:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  5 /job:worker/replica:0/task:0/device:TPU:5: <tf.Variable 'dense/kernel/replica_5:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  6 /job:worker/replica:0/task:0/device:TPU:6: <tf.Variable 'dense/kernel/replica_6:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  7 /job:worker/replica:0/task:0/device:TPU:7: <tf.Variable 'dense/kernel/replica_7:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>
}]
[TPUMirroredVariable:{
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense_1/kernel:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  1 /job:worker/replica:0/task:0/device:TPU:1: <tf.Variable 'dense_1/kernel/replica_1:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  2 /job:worker/replica:0/task:0/device:TPU:2: <tf.Variable 'dense_1/kernel/replica_2:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  3 /job:worker/replica:0/task:0/device:TPU:3: <tf.Variable 'dense_1/kernel/replica_3:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  4 /job:worker/replica:0/task:0/device:TPU:4: <tf.Variable 'dense_1/kernel/replica_4:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  5 /job:worker/replica:0/task:0/device:TPU:5: <tf.Variable 'dense_1/kernel/replica_5:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  6 /job:worker/replica:0/task:0/device:TPU:6: <tf.Variable 'dense_1/kernel/replica_6:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,
  7 /job:worker/replica:0/task:0/device:TPU:7: <tf.Variable 'dense_1/kernel/replica_7:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>
}]
Update only B
[TPUMirroredVariable:{
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense/kernel:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  1 /job:worker/replica:0/task:0/device:TPU:1: <tf.Variable 'dense/kernel/replica_1:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  2 /job:worker/replica:0/task:0/device:TPU:2: <tf.Variable 'dense/kernel/replica_2:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  3 /job:worker/replica:0/task:0/device:TPU:3: <tf.Variable 'dense/kernel/replica_3:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  4 /job:worker/replica:0/task:0/device:TPU:4: <tf.Variable 'dense/kernel/replica_4:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  5 /job:worker/replica:0/task:0/device:TPU:5: <tf.Variable 'dense/kernel/replica_5:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  6 /job:worker/replica:0/task:0/device:TPU:6: <tf.Variable 'dense/kernel/replica_6:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  7 /job:worker/replica:0/task:0/device:TPU:7: <tf.Variable 'dense/kernel/replica_7:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>
}]
[TPUMirroredVariable:{
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense_1/kernel:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.219872 ],
       [ 1.3811986]], dtype=float32)>,
  1 /job:worker/replica:0/task:0/device:TPU:1: <tf.Variable 'dense_1/kernel/replica_1:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.219872 ],
       [ 1.3811986]], dtype=float32)>,
  2 /job:worker/replica:0/task:0/device:TPU:2: <tf.Variable 'dense_1/kernel/replica_2:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.219872 ],
       [ 1.3811986]], dtype=float32)>,
  3 /job:worker/replica:0/task:0/device:TPU:3: <tf.Variable 'dense_1/kernel/replica_3:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.219872 ],
       [ 1.3811986]], dtype=float32)>,
  4 /job:worker/replica:0/task:0/device:TPU:4: <tf.Variable 'dense_1/kernel/replica_4:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.219872 ],
       [ 1.3811986]], dtype=float32)>,
  5 /job:worker/replica:0/task:0/device:TPU:5: <tf.Variable 'dense_1/kernel/replica_5:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.219872 ],
       [ 1.3811986]], dtype=float32)>,
  6 /job:worker/replica:0/task:0/device:TPU:6: <tf.Variable 'dense_1/kernel/replica_6:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.219872 ],
       [ 1.3811986]], dtype=float32)>,
  7 /job:worker/replica:0/task:0/device:TPU:7: <tf.Variable 'dense_1/kernel/replica_7:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.219872 ],
       [ 1.3811986]], dtype=float32)>
}

TPUの8個分のデバイスの出力が出てちょっとわかりづらいです（Distribute Trainingのため）が、0番目のデバイスだけ着目してみましょう。

Before update
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense/kernel:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.31612158, -0.6931683 ]], dtype=float32)>,
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense_1/kernel:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,

Update only A 
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense/kernel:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense_1/kernel:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.2208728],
       [ 1.3821994]], dtype=float32)>,

Update only B
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense/kernel:0' shape=(1, 2) dtype=float32, numpy=array([[ 0.3151208, -0.6921675]], dtype=float32)>,
  0 /job:worker/replica:0/task:0/device:TPU:0: <tf.Variable 'dense_1/kernel:0' shape=(2, 1) dtype=float32, numpy=
array([[-1.219872 ],
       [ 1.3811986]], dtype=float32)>,

Update only A では、モデルAの係数のみ変わっており、モデルBの係数は変わっていません。

Update only B では、モデルBの係数のみ変わっており、モデルAの係数は変わっていません。

→うまくいってる

この結果を踏まえると、TPUでもいい感じにGANは実装できそうです。

サークルからのお知らせ

Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」では、技術書・その他の本のごったまぜ通販をしています！　大人気のモザイク除去から北の駅巡りまでごちゃごちゃ感をお楽しみください！

技術書コーナー

(1) モザイク除去から学ぶ最先端のディープラーニング

技術書典8の新刊『モザイク除去から学ぶ　最先端のディープラーニング』好評通販中（A4・195ページ）です！　Boothから試し読みできます！

これが君の望んでいる、「ディープラーニング－ロマン－」

誰もが夢見る「モザイク除去」から機械学習を学ぼう。ご好評につき鋭意増刷中！

Booth：https://t.co/d2oXBf0kQ6
とら：https://t.co/5P0HJ7r8sL
試し読み：https://t.co/1z2XScVNsz #技術書典 pic.twitter.com/tttO6VnTXG
— こしあん@技術書典8 Day2 あ03　委託 (@js_koshian2) March 22, 2020

Booth（物理、電子版）：https://koshian2.booth.pm/items/1835219
とらのあな（物理版）：https://ec.toranoana.shop/tora/ec/item/040030818462/
メロンブックス（物理版、電子版とセット）https://www.melonbooks.co.jp/detail/detail.php?product_id=670234
感想・情報まとめ・質問用：https://github.com/koshian2/MosaicDeeplearningBook

(2) Shikoan's ML Blog

このブログが電子書籍になりました！　1冊オムニバス形式で200P超の大ボリューム。どこから読んでもOKなので、ぜひ通勤・通学のお供にどうぞ。

累計100万PV超の人気ブログが待望の電子化！　

・どこから読んでも、何巻から読んでもOK

・短い時間でもすぐ読めて便利な「しおり」機能つき！

・1冊にたっぷり30本収録。1本あたり18.3円！
・文庫本感覚でどうぞ！

★ Vol.1 ★https://t.co/sDZrHdnotq

★ Vol.2 ★https://t.co/URKqcVWUZs pic.twitter.com/ZYxiUGbCkV
— こしあん@11/14-15　アキバ同人祭「ラ29」「日高本線　車なし全駅めぐり」 (@js_koshian2) November 2, 2020

Vol.1 https://koshian2.booth.pm/items/2169975
Vol.2 https://koshian2.booth.pm/items/2170021

北の鉄道コーナー

Vol.1:日高本線　車なし全駅巡り

2021年に大半が廃線の日高本線を、公共交通機関だけ、マイカー禁止で全駅訪問した記録。最終的に頼れるのは「徒歩」！

数少ない代行バス、複雑怪奇な路線バス、限界徒歩の先にある日高のグルメとは？　でかい・厚いのフルカラー本。シリーズ記念第一弾！

★ Booth ★https://t.co/hbDgkNgsey pic.twitter.com/3rCTuNyr1b
— こしあん@11/14-15　アキバ同人祭「ラ29」「日高本線　車なし全駅めぐり」 (@js_koshian2) November 2, 2020

Booth（物理＋電子＋おまけ壁紙128枚つき）：https://koshian2.booth.pm/items/2334590

Pocket

Tags:DeepLearning, GAN, TPU, 小ネタ

書籍情報

コード

結果

Add a Comment