TPUを仕組み理解からTensorFlowで実装まで
TPUって何?
見ていると気になることはいっぱいある
上げればキリがないものだが、昔のメモをさグッたらTPUという文字が出てきたのでこの際に調べておこう。
Tensor Processing Unit
最も事細かにかかなので、あしからず。
仕組みの説明はGoogleCloudの説明がしっくりきた。
まず、CPU・GPUの復習
- CPU
- ALUに何でもやらせるから柔軟性は高い。
- でもキャッシュ読み込み→命令処理→書き込みまでが時間とエネルギーの無駄
- パイプライン処理で頑張っても、結局たかが知れてる
- GPU
- 1個のプロセッサに大量のALUを入れたもの。単純な命令を「数」で倒せる。
- 数の分だけ高速化するが、電力を食う。
- CPUと仕組みは同じなので、中間でキャッシュを経由するため時間とエネルギーの無駄
- この呪縛をノイマンボトルネックというらしい。この名は初めて知った
Googleはうまいことを考えた。GPU・CPUの欠点を克服するには、構造をかえようと
キャッシュに行く前にどんどん足してしまえばよい!
Google公式のGIFがよいね
データを左から列方向に展開して→に進めていく
パラメータを下から行方向に展開して↑に進めていく
NNっていうのは入力データに対し、パラメータを1対1(y=ax+b)でつなげるものなので、
上2図のように流していけば、否が応でも演算器にぶち当たり計算せざるを得ない。
出口である右端では最終的に下みたいな行列演算(matmul)ができていることになる
感動した。頭いい。
このキャッシュ通さない作戦をシストリックアレイって言うらしい。
また、NNの計算過程では精度はそこまで重要ではないらしく、
量子化を行って一般的な32/64bit演算から8bit演算の仕組みを採用したため、実質4倍高速化(=省エネ)
この精度周りの下りは結構アツいらしい
https://ascii.jp/elem/000/004/014/4014066/
ムーアの法則って呪縛に聞こえてきた…
GPUとCPUのお互いいいところがあるように、TPUにも向き不向きがある。
Googleから直接拝借する。
・CPU
https://cloud.google.com/tpu/docs/tpus?hl=ja
最大限の柔軟性を必要とする迅速なプロトタイピング
トレーニングに時間がかからない単純なモデル
実際のバッチサイズが小さい小規模なモデル
C++ で記述されたカスタム TensorFlow 演算が多くを占めるモデル
ホストシステムの使用可能な I/O またはネットワーク帯域幅によって制限が課せられるモデル
・GPU
ソースが存在しないモデルまたはソースを変更するのが煩雑すぎるモデル
CPU 上で少なくとも部分的に実行しなければならない多数のカスタム TensorFlow 演算を使用するモデル
Cloud TPU で利用できない TensorFlow 演算を使用するモデル(利用可能な TensorFlow 演算のリストをご覧ください)
実際のバッチサイズが大きい中〜大規模なモデル
・TPU
行列計算が多くを占めるモデル
メインのトレーニング ループ内にカスタム TensorFlow 演算がないモデル
トレーニングに数週間または数か月かかるモデル
実際のバッチサイズが非常に大きい非常に大規模なモデル
GPUはある程度複雑な計算でもある程度応用(適応?)が効いたが、TPUは一切それが効かない。
その代わり、単純な行列演算の塊なら必殺!といった次第だろう
ColabとかKaggleとかでは使えるのでつかってみた
実装してみます。とは言っても他人コードを拝借する。一部加えつつ、、
1年前にkaggaleでTPUつかおうぜ!みたいなのを思い出したので、それを題材にします。
お花分類ですね
こちらに書いてあるTPUの使い方はとてもわかりやすいと思います(あくまで理解のうえでは)
SAMPLEコードにめもを追加していきます。
1 2 3 4 5 6 7 8 9 10 11 12 13 |
import re import tensorflow as tf import tensorflow_hub as hub import numpy as np from matplotlib import pyplot as plt print("Tensorflow version " + tf.__version__) # Prefetch用にどの程度先取りしておくか自動調整 # これもいつもどおりでOK # https://tensorflow.classcat.com/2019/03/23/tf20-alpha-guide-data-performance/ AUTO = tf.data.experimental.AUTOTUNE # Kaggleのデータ拝借 from kaggle_datasets import KaggleDatasets |
TPUとGPUを切り替える部分。
各TPUを制御するcluster_resolverを使ってTPUをネットワークから探索。
各TPU間の同期方法を定義するstrategyは必須。strategyはあくまでインスタンスなので、のちほど定義する
TPUStrategyの場合、アクセラレーターで分散する際に変数のミラーリングがされ、分散学習の重みを集約→重みを更新(同期)→共有する。この仕組みをall-reduceっていう。
TPUは使えない場合はGPUの同じミラーリング戦略であるMirroredStrategyをつかう
以下、日経 xTECHを拝借
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
try: tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect() # TPU detection strategy = tf.distribute.TPUStrategy(tpu) print("Detect: TPU") except ValueError: # detect GPUs strategy = tf.distribute.MirroredStrategy() # for GPU or multi-GPU machines #strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU(インストールで入れた際に) #strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy() # for clusters of multi-GPU machines print("Detect: GPU") print("Number of accelerators: ", strategy.num_replicas_in_sync) # tf.distribute.Strategyは、トレーニングを複数のGPU、複数のマシン、またはTPUに分散するためのTensorFlowAPI # https://tensorflow.classcat.com/2019/03/21/tf20-alpha-guide-distribute-strategy/ |
まずはデータを撮ってくる。
バッチサイズが大事で、ここは8の倍数(GoogleTPUv3.8の場合8bit量子化しているので)にしないと行けない。
これは後で使うSteps_per_executionも同じことが言える。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
""" Define image info """ # Dataset: https://www.kaggle.com/mgornergoogle/five-flowers # TFRecord ファイル群 EPOCHS = 12 IMAGE_SIZE = [331, 331] # available image sizes # train/test/val FLOWERS_DATASETS = { 192: GCS_PATH + '/tfrecords-jpeg-192x192/*.tfrec', 224: GCS_PATH + '/tfrecords-jpeg-224x224/*.tfrec', 331: GCS_PATH + '/tfrecords-jpeg-331x331/*.tfrec', 512: GCS_PATH + '/tfrecords-jpeg-512x512/*.tfrec' } # Labels CLASSES = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips'] assert IMAGE_SIZE[0] == IMAGE_SIZE[1], "only square images are supported" assert IMAGE_SIZE[0] in FLOWERS_DATASETS, "this image size is not supported" # strategy.num_replicas_in_sync: #TPU or #GPU: レプリカを各GPU/TPUで作成 BATCH_SIZE = 16 * strategy.num_replicas_in_sync |
学習率を途中で変えるらしいです。TPUに限定してなぜカスタマイズするのかは不明。
ただ自分自身で変更するのは初めて。よい知見。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
""" learning rate scheduler """ LR_START = 0.00001 LR_MAX = 0.00005 * strategy.num_replicas_in_sync LR_MIN = 0.00001 LR_RAMPUP_EPOCHS = 5 LR_SUSTAIN_EPOCHS = 0 LR_EXP_DECAY = .8 def lrfn(epoch): if epoch < LR_RAMPUP_EPOCHS: lr = (LR_MAX - LR_START) / LR_RAMPUP_EPOCHS * epoch + LR_START elif epoch < LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS: lr = LR_MAX else: # 以降指数的に減衰 lr = (LR_MAX - LR_MIN) * LR_EXP_DECAY**(epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS) + LR_MIN return lr lr_callback = tf.keras.callbacks.LearningRateScheduler(lrfn, verbose=True) rng = [i for i in range(EPOCHS)] y = [lrfn(x) for x in rng] plt.plot(rng, y) print("Learning rate schedule: {:.3g} to {:.3g} to {:.3g}".format(y[0], max(y), y[-1])) |
この辺は画像を撮ってくるだけ。TFRecordなので、それに応じて読み込む。
ここは全て借り物。おパクり。
重要な点は最初に設定した”AUTO”。
TPU・GPUはPrefetchが速度のボトルネックなのでここは大事。それ以外はデータよりけりだと思う
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
""" Data Loading Funcs 全て借り物関数 """ #TPUはサイズに気をつけないといけないらしい #https://cloud.google.com/tpu/docs/performance-guide#tf-fcts # https://qiita.com/ohtaman/items/325615358f008a4005e2 def count_data_items(filenames): # the number of data items is written in the name of the .tfrec files, i.e. flowers00-230.tfrec = 230 data items n = [int(re.compile(r"-([0-9]*)\.").search(filename).group(1)) for filename in filenames] return np.sum(n) def read_tfrecord(example): features = { "image": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring "class": tf.io.FixedLenFeature([], tf.int64), # shape [] means scalar "one_hot_class": tf.io.VarLenFeature(tf.float32),# Load variable length feature } example = tf.io.parse_single_example(example, features) image = tf.image.decode_jpeg(example['image'], channels=3) # pixel format uint8 [0,255] range class_label = tf.cast(example['class'], tf.int32) # not used one_hot_class = tf.sparse.to_dense(example['one_hot_class']) one_hot_class = tf.reshape(one_hot_class, [5]) return image, one_hot_class def force_image_sizes(dataset, image_size): # explicit size needed for TPU reshape_images = lambda image, label: (tf.reshape(image, [*image_size, 3]), label) dataset = dataset.map(reshape_images, num_parallel_calls=AUTO) return dataset def load_dataset(filenames): # Read from TFRecords. For optimal performance, reading from multiple files at once and # disregarding data order. Order does not matter since we will be shuffling the data anyway. ignore_order = tf.data.Options() ignore_order.experimental_deterministic = False dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO) # automatically interleaves reads from multiple files dataset = dataset.with_options(ignore_order) # uses data as soon as it streams in, rather than in its original order dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTO) dataset = force_image_sizes(dataset, IMAGE_SIZE) return dataset def data_augment(image, one_hot_class): # data augmentation. Thanks to the dataset.prefetch(AUTO) statement in the next function (below), # this happens essentially for free on TPU. Data pipeline code is executed on the "CPU" part # of the TPU while the TPU itself is computing gradients. image = tf.image.random_flip_left_right(image) image = tf.image.random_saturation(image, 0, 2) return image, one_hot_class def get_training_dataset(): dataset = load_dataset(TRAINING_FILENAMES) dataset = dataset.map(data_augment, num_parallel_calls=AUTO) dataset = dataset.repeat() dataset = dataset.shuffle(2048) dataset = dataset.batch(BATCH_SIZE) dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size) return dataset def get_validation_dataset(): dataset = load_dataset(VALIDATION_FILENAMES) dataset = dataset.batch(BATCH_SIZE) dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size) return dataset |
よしなに分割。
1 2 3 4 5 6 7 8 9 10 11 12 13 |
""" Split Data """ gcs_pattern = FLOWERS_DATASETS[IMAGE_SIZE[0]] validation_split = 0.19 filenames = tf.io.gfile.glob(gcs_pattern) split = len(filenames) - int(len(filenames) * validation_split) TRAINING_FILENAMES = filenames[:split] VALIDATION_FILENAMES = filenames[split:] TRAIN_STEPS = count_data_items(TRAINING_FILENAMES) // BATCH_SIZE VALIDATION_STEPS = -(-count_data_items(VALIDATION_FILENAMES) // BATCH_SIZE) # The "-(-//)" trick rounds up instead of down :-) print("TRAINING IMAGES: ", count_data_items(TRAINING_FILENAMES), ", STEPS PER EPOCH: ", TRAIN_STEPS) print("VALIDATION IMAGES: ", count_data_items(VALIDATION_FILENAMES)) |
モデル作成。
strategy.scope()でくくることで、先程のインスタンスがようやく登場。
この中の処理のみがTPUで共有される。つまり、モデル部分だけが共有される。
今回の例ではSAMPLEに従って、xceptionを採用。下記の通り他のモデルを使いたいならテキトウに変更
steps_per_execution=8は、1callでなんバッチ読み込むかを定義する。上の”AUTO”と同じで高速化に貢献
単純なギモンとしてGPUでも同じこと言えるよね?と思ったが、GPUは関係ないらしい。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
# xceptionを利用 # https://note.nkmk.me/python-tensorflow-keras-applications-pretrained-models/ # VGGとかこの部分を変えればよい。 # ベつに自作してもよい。普通の画像分類 with strategy.scope(): """ Distributed Execution """ # xception用のn前処理実行 img_adjust_layer = tf.keras.layers.Lambda( lambda data: tf.keras.applications.xception.preprocess_input( tf.cast(data, tf.float32)), input_shape=[*IMAGE_SIZE, 3]) # xceptionモデルの組み込み pretrained_model = tf.keras.applications.Xception(weights='imagenet', include_top=False) #学習モード pretrained_model.trainable = True model = tf.keras.Sequential([ img_adjust_layer, pretrained_model, # GAP>Flat>Denceで終端処理 tf.keras.layers.GlobalAveragePooling2D(), tf.keras.layers.Flatten(), tf.keras.layers.Dense(5, activation='softmax') ]) # steps_per_execution: 1回のrunで何バッチを呼ぶか? # TPUまたは小さいモデルなら効果あり(50%高速化)。 # GPUがなぜ効果内のか不明。 #https://keras.io/api/models/model_training_apis/ # https://qiita.com/T-STAR/items/e2998d4c22c882039ffb model.compile( optimizer='adam', loss = 'categorical_crossentropy', metrics=['accuracy'], steps_per_execution=8 ) model.summary() |
学習。もはやGPUのときと変わらない。
1 2 3 4 5 6 |
history = model.fit(get_training_dataset(), steps_per_epoch=TRAIN_STEPS, epochs=EPOCHS, validation_data=get_validation_dataset(), validation_steps=VALIDATION_STEPS, callbacks=[lr_callback]) final_accuracy = history.history["val_accuracy"][-5:] print("FINAL ACCURACY MEAN-5: ", np.mean(final_accuracy)) |
このあと保存と読み込みがあるが、GPUと読み込み方法は異なるものの、余り本質ではないんので省略。
いまはTenswflow liteというものがあって、いわゆるエッジコンピューティング用に特化されたフレームワークがあるけど、そこでも活用できるみたい。
電子工作でおなじみのswitchScienceから出ている。今度遊んでみたい。1万するけどwww
https://www.switch-science.com/catalog/5817/
gmoがなんかやっていた。
感想
- TPUナニそれすごそう!という印象だったが、案外GPUと実装の手間がほぼ変わらない。クラウドならかんたんに使えそう
- 久しぶりにCSチックな言葉が出てきたので興奮した。大学院のときよりも
補足
- 課金とか実際のGCPでの利用は以下が使えそう
- 今回のコード