1. 概要
TPU は非常に高速です。トレーニング データのストリームも、トレーニングのスピードについていく必要があります。このラボでは、tf.data.Dataset API を使用して GCS からデータを読み込み、TPU にフィードする方法を学びます。
このラボは、「Keras on TPU」シリーズのパート 1 です。これらの手順は、次の順序で実行することも、個別に行うこともできます。
- [このラボ] TPU スピードのデータ パイプライン(tf.data.Dataset と TFRecords)
- 転移学習を使用した最初の Keras モデル
- Keras と TPU を使用した畳み込みニューラル ネットワーク
- Keras と TPU での最新の convnets、squeezenet、Xception

学習内容
- tf.data.Dataset API を使用してトレーニング データを読み込む
- TFRecord 形式を使用して、GCS から効率的にトレーニング データを読み込む
フィードバック
この Codelab で問題が見つかった場合は、お知らせください。フィードバックは GitHub の問題 [フィードバック リンク] から送信できます。
2. Google Colaboratory クイック スタート
このラボでは Google Collaboratory を使用するため、ユーザー側での設定は不要です。Colaboratory は、教育目的のオンライン ノートブック プラットフォームです。CPU、GPU、TPU のトレーニングを無料で提供しています。

このサンプル ノートブックを開いて、いくつかのセルを実行すると、Colaboratory の使い方を理解できます。
TPU バックエンドを選択する

Colab のメニューで、[ランタイム] > [ランタイムのタイプを変更] を選択し、[TPU] を選択します。このコードラボでは、ハードウェア アクセラレータによるトレーニングをサポートする強力な TPU(Tensor Processing Unit)を使用します。ランタイムへの接続は、初回実行時に自動的に行われます。右上にある [接続] ボタンを使用することもできます。
ノートブックの実行

セルを 1 つずつ実行するには、セルをクリックして Shift+Enter キーを押します。[ランタイム] > [すべて実行] を選択して、ノートブック全体を実行することもできます。
目次

すべてのノートブックに目次があります。左側の黒い矢印を使用して開くことができます。
非表示のセル

一部のセルにはタイトルのみが表示されます。これは Colab 固有のノートブック機能です。ダブルクリックすると内部のコードを確認できますが、通常はあまり興味深いものではありません。通常はサポート関数または可視化関数。内部の関数を定義するには、これらのセルを実行する必要があります。
認証

承認済みアカウントで認証すれば、Colab からプライベート Google Cloud Storage バケットにアクセスできます。上記のコード スニペットは、認証プロセスをトリガーします。
3. [INFO] Tensor Processing Unit(TPU)とは
概要

Keras で TPU 上でモデルをトレーニングするコード(TPU が使用できない場合は GPU または CPU にフォールバックします)。
try: # detect TPUs
tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
strategy = tf.distribute.TPUStrategy(tpu)
except ValueError: # detect GPUs
strategy = tf.distribute.MirroredStrategy() # for CPU/GPU or multi-GPU machines
# use TPUStrategy scope to define model
with strategy.scope():
model = tf.keras.Sequential( ... )
model.compile( ... )
# train model normally on a tf.data.Dataset
model.fit(training_dataset, epochs=EPOCHS, steps_per_epoch=...)
今日は、TPU を使用して、インタラクティブな速度(トレーニング実行あたり数分)で花の分類子を構築して最適化します。

TPU を使用する理由
最新の GPU は、プログラマブルな「コア」を中心に構成されています。これは、3D レンダリング、ディープ ラーニング、物理シミュレーションなど、さまざまなタスクを処理できる非常に柔軟なアーキテクチャです。一方、TPU は、従来のベクトル プロセッサと専用の行列乗算ユニットを組み合わせたもので、ニューラル ネットワークなど、大規模な行列乗算が支配的なタスクに優れています。

図: 密ニューラル ネットワーク レイヤを行列乗算として表した図。8 枚の画像のバッチがニューラル ネットワークで一度に処理されます。1 行 x 列の乗算を実行して、画像内のすべてのピクセル値の加重和が実際に計算されていることを確認してください。畳み込みレイヤも行列乗算として表現できますが、少し複雑になります( 説明はこちらのセクション 1 をご覧ください)。
ハードウェア
MXU と VPU
TPU v2 コアは、行列乗算を実行する Matrix Multiply Unit(MXU)と、活性化やソフトマックスなどの他のすべてのタスク用の Vector Processing Unit(VPU)で構成されています。VPU は float32 と int32 の計算を処理します。一方、MXU は混合精度 16 ~ 32 ビットの浮動小数点形式で動作します。

混合精度の浮動小数点と bfloat16
MXU は、bfloat16 入力と float32 出力を使用して行列乗算を計算します。中間累積は float32 精度で実行されます。

ニューラル ネットワークのトレーニングは通常、浮動小数点精度の低下によって生じるノイズの影響を受けません。ノイズがオプティマイザーの収束に役立つ場合もあります。16 ビット浮動小数点精度は、従来、計算の高速化に使用されてきましたが、float16 形式と float32 形式では範囲が大きく異なります。精度を float32 から float16 に下げると、通常はオーバーフローとアンダーフローが発生します。ソリューションは存在しますが、通常は float16 を機能させるために追加の作業が必要です。
そのため、Google は TPU に bfloat16 形式を導入しました。bfloat16 は、float32 とまったく同じ指数ビットと範囲を持つ切り捨てられた float32 です。また、TPU は bfloat16 入力と float32 出力を使用して混合精度で行列乗算を計算するため、通常は、精度を下げたことによるパフォーマンスの向上を利用するためにコードを変更する必要はありません。
シストリック アレイ
MXU は、データ要素がハードウェア コンピューティング ユニットの配列を流れる「シストリック アレイ」アーキテクチャを使用して、ハードウェアで行列乗算を実装します。(医学では、収縮期とは心臓の収縮と血流を指しますが、ここではデータの流れを指します)。
行列乗算の基本要素は、一方の行列の行と他方の行列の列の内積です(このセクションの上部の図を参照)。行列乗算 Y=X*W の場合、結果の 1 つの要素は次のようになります。
Y[2,0] = X[2,0]*W[0,0] + X[2,1]*W[1,0] + X[2,2]*W[2,0] + ... + X[2,n]*W[n,0]
GPU では、このドット積を GPU の「コア」にプログラムし、利用可能なすべての「コア」で並列実行して、結果の行列のすべての値を一度に計算しようとします。結果の行列が 128x128 の大きさの場合、128x128=16K 個の「コア」が必要になりますが、通常これは不可能です。最大の GPU には約 4,000 個のコアがあります。一方、TPU は MXU のコンピューティング ユニットに最小限のハードウェア(bfloat16 x bfloat16 => float32 乗算アキュムレータのみ)を使用します。これらは非常に小さいため、TPU は 128x128 MXU に 16 K 個を実装し、この行列乗算を一度に処理できます。

図: MXU シストリック アレイ。コンピューティング要素は乗算アキュムレータです。1 つの行列の値が配列に読み込まれます(赤い点)。もう一方の行列の値は配列を通過します(灰色の点)。縦線は値を上に伝播します。水平線は部分和を伝播します。配列をデータが流れるときに、右側から行列乗算の結果が得られることを確認するのは、ユーザーの演習として残されています。
また、MXU でドット積が計算されている間、中間合計は隣接するコンピューティング ユニット間を流れるだけです。メモリやレジスタ ファイルに保存して取得する必要はありません。その結果、TPU シストリック アレイ アーキテクチャは、行列乗算の計算において、GPU よりも密度と電力の面で大きな利点があり、速度の面でも無視できない利点があります。
Cloud TPU
Google Cloud Platform で「Cloud TPU v2」をリクエストすると、PCI 接続の TPU ボードを備えた仮想マシン(VM)が取得されます。TPU ボードには、4 つのデュアルコア TPU チップがあります。各 TPU コアには、VPU(ベクトル処理ユニット)と 128x128 MXU(マトリックス乗算ユニット)があります。この「Cloud TPU」は通常、ネットワークを介してリクエストした VM に接続されます。全体像は次のようになります。

図: ネットワーク接続された「Cloud TPU」アクセラレータを備えた VM。「Cloud TPU」自体は、PCI 接続の TPU ボードと、そのボードに搭載された 4 つのデュアルコア TPU チップを備えた VM で構成されています。
TPU Pod
Google のデータセンターでは、TPU はハイ パフォーマンス コンピューティング(HPC)相互接続に接続されており、1 つの非常に大きなアクセラレータとして認識されることがあります。Google はこれらを Pod と呼び、最大 512 個の TPU v2 コアまたは 2,048 個の TPU v3 コアを包含できます。

図: TPU v3 Pod。HPC 相互接続を介して接続された TPU ボードとラック。
トレーニング中、勾配は all-reduce アルゴリズム(all-reduce の詳細はこちら)を使用して TPU コア間で交換されます。トレーニング対象のモデルは、大きなバッチサイズでトレーニングすることでハードウェアを活用できます。

図: Google TPU の 2 次元トーラス メッシュ HPC ネットワークで all-reduce アルゴリズムを使用してトレーニング中に勾配を同期する。
ソフトウェア
大規模なバッチサイズのトレーニング
TPU の理想的なバッチサイズは TPU コアあたり 128 個のデータ項目ですが、ハードウェアは TPU コアあたり 8 個のデータ項目から高い使用率を示します。1 つの Cloud TPU には 8 個のコアがあります。
この Codelab では、Keras API を使用します。Keras で指定するバッチは、TPU 全体のグローバル バッチサイズです。バッチは自動的に 8 つに分割され、TPU の 8 個のコアで実行されます。

パフォーマンスに関するその他のヒントについては、TPU パフォーマンス ガイドをご覧ください。バッチサイズが非常に大きい場合は、一部のモデルで特別な配慮が必要になることがあります。詳しくは、LARSOptimizer をご覧ください。
仕組み: XLA
Tensorflow プログラムは計算グラフを定義します。TPU は Python コードを直接実行するのではなく、TensorFlow プログラムで定義された計算グラフを実行します。内部では、XLA(高速線形代数コンパイラ)と呼ばれるコンパイラが、計算ノードの Tensorflow グラフを TPU マシンコードに変換します。このコンパイラは、コードとメモリ レイアウトに対して多くの高度な最適化も実行します。コンパイルは、作業が TPU に送信されると自動的に行われます。ビルドチェーンに XLA を明示的に含める必要はありません。

図: TPU で実行するには、TensorFlow プログラムで定義された計算グラフをまず XLA(アクセラレータ線形代数コンパイラ)表現に変換し、次に XLA によって TPU マシンコードにコンパイルします。
Keras で TPU を使用する
Tensorflow 2.1 以降、Keras API を介して TPU がサポートされています。Keras のサポートは、TPU と TPU Pod で機能します。TPU、GPU、CPU で動作する例を次に示します。
try: # detect TPUs
tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
strategy = tf.distribute.TPUStrategy(tpu)
except ValueError: # detect GPUs
strategy = tf.distribute.MirroredStrategy() # for CPU/GPU or multi-GPU machines
# use TPUStrategy scope to define model
with strategy.scope():
model = tf.keras.Sequential( ... )
model.compile( ... )
# train model normally on a tf.data.Dataset
model.fit(training_dataset, epochs=EPOCHS, steps_per_epoch=...)
このコード スニペットでは:
TPUClusterResolver().connect()は、ネットワーク上の TPU を検出します。ほとんどの Google Cloud システム(AI Platform ジョブ、Colaboratory、Kubeflow、ctpu up ユーティリティで作成された Deep Learning VM)でパラメータなしで動作します。これらのシステムは、TPU_NAME 環境変数のおかげで TPU の場所を認識しています。TPU を手動で作成する場合は、使用する VM で TPU_NAME 環境変数を設定するか、明示的なパラメータを指定してTPUClusterResolverを呼び出します。TPUClusterResolver(tp_uname, zone, project)TPUStrategyは、分布と「オールリデュース」勾配同期アルゴリズムを実装する部分です。- この戦略はスコープを通じて適用されます。モデルは strategy scope() 内で定義する必要があります。
tpu_model.fit関数は、TPU トレーニングの入力として tf.data.Dataset オブジェクトを想定しています。
一般的な TPU の移植タスク
- Tensorflow モデルでデータを読み込む方法は多数ありますが、TPU では
tf.data.DatasetAPI を使用する必要があります。 - TPU は非常に高速であるため、TPU で実行すると、データの取り込みがボトルネックになることがよくあります。TPU パフォーマンス ガイドには、データ ボトルネックを検出するために使用できるツールや、その他のパフォーマンスに関するヒントが記載されています。
- int8 または int16 の数値は int32 として扱われます。TPU には、32 ビット未満で動作する整数ハードウェアはありません。
- 一部の TensorFlow オペレーションはサポートされていません。リストはこちらをご覧ください。この制限はトレーニング コード(モデルのフォワード パスとバックワード パス)にのみ適用されます。データ入力パイプラインは CPU で実行されるため、Tensorflow のすべてのオペレーションをデータ入力パイプラインで使用できます。
- TPU では
tf.py_funcはサポートされていません。
4. データを読み込んでいます

ここでは、花の画像データセットを使用します。目標は、それらを 5 種類の花に分類する方法を学習することです。データの読み込みは tf.data.Dataset API を使用して行われます。まず、API について理解しましょう。
実践的
次のノートブックを開き、セルを実行(Shift+Enter)して、「WORK REQUIRED」ラベルが表示されている手順に沿って操作してください。
Fun with tf.data.Dataset (playground).ipynb
その他の情報
「flowers」データセットについて
データセットは 5 つのフォルダに整理されています。各フォルダには 1 種類の花が含まれています。フォルダの名前は、sunflowers、daisy、dandelion、tulips、roses です。データは Google Cloud Storage の公開バケットでホストされています。抜粋:
gs://flowers-public/sunflowers/5139971615_434ff8ed8b_n.jpg
gs://flowers-public/daisy/8094774544_35465c1c64.jpg
gs://flowers-public/sunflowers/9309473873_9d62b9082e.jpg
gs://flowers-public/dandelion/19551343954_83bb52f310_m.jpg
gs://flowers-public/dandelion/14199664556_188b37e51e.jpg
gs://flowers-public/tulips/4290566894_c7f061583d_m.jpg
gs://flowers-public/roses/3065719996_c16ecd5551.jpg
gs://flowers-public/dandelion/8168031302_6e36f39d87.jpg
gs://flowers-public/sunflowers/9564240106_0577e919da_n.jpg
gs://flowers-public/daisy/14167543177_cd36b54ac6_n.jpg
tf.data.Dataset を使用する理由
Keras と TensorFlow は、すべてのトレーニング関数と評価関数でデータセットを受け入れます。データセットにデータを読み込むと、API はニューラル ネットワークのトレーニング データに役立つ一般的な機能をすべて提供します。
dataset = ... # load something (see below)
dataset = dataset.shuffle(1000) # shuffle the dataset with a buffer of 1000
dataset = dataset.cache() # cache the dataset in RAM or on disk
dataset = dataset.repeat() # repeat the dataset indefinitely
dataset = dataset.batch(128) # batch data elements together in batches of 128
AUTOTUNE = tf.data.AUTOTUNE
dataset = dataset.prefetch(AUTOTUNE) # prefetch next batch(es) while training
パフォーマンスのヒントとデータセットのベスト プラクティスについては、こちらの記事をご覧ください。リファレンス ドキュメントはこちらをご覧ください。
tf.data.Dataset の基本
通常、データは複数のファイル(ここでは画像)で提供されます。ファイル名のデータセットは、次の呼び出しで作成できます。
filenames_dataset = tf.data.Dataset.list_files('gs://flowers-public/*/*.jpg')
# The parameter is a "glob" pattern that supports the * and ? wildcards.
次に、各ファイル名に関数を「マッピング」します。通常、この関数はファイルを読み込んでデコードし、メモリ内の実際のデータに変換します。
def decode_jpeg(filename):
bits = tf.io.read_file(filename)
image = tf.io.decode_jpeg(bits)
return image
image_dataset = filenames_dataset.map(decode_jpeg)
# this is now a dataset of decoded images (uint8 RGB format)
データセットを反復処理するには:
for data in my_dataset:
print(data)
タプルのデータセット
教師あり学習では、通常、トレーニング データセットはトレーニング データと正解のペアで構成されます。これを可能にするため、デコード関数はタプルを返すことができます。タプルのデータセットが作成され、そのデータセットを反復処理するとタプルが返されます。返される値は、モデルで使用できる Tensorflow テンソルです。.numpy() を呼び出すと、未加工の値を確認できます。
def decode_jpeg_and_label(filename):
bits = tf.read_file(filename)
image = tf.io.decode_jpeg(bits)
label = ... # extract flower name from folder name
return image, label
image_dataset = filenames_dataset.map(decode_jpeg_and_label)
# this is now a dataset of (image, label) pairs
for image, label in dataset:
print(image.numpy().shape, label.numpy())
結論:画像を 1 つずつ読み込むのは遅い
このデータセットを反復処理すると、1 秒あたり 1 ~ 2 枚の画像を読み込むことができることがわかります。遅すぎます。トレーニングに使用するハードウェア アクセラレータは、このレートの何倍ものレートを維持できます。次のセクションでは、この方法について説明します。
ソリューション
ソリューション ノートブックは次のとおりです。行き詰まった場合は、こちらをご利用ください。
Fun with tf.data.Dataset (solution).ipynb
学習した内容
- 🤔 tf.data.Dataset.list_files
- 🤔 tf.data.Dataset.map
- 🤔 タプルのデータセット
- 😀 データセットの反復処理
このチェックリストを頭の中で確認してください。
5. データを高速で読み込む
このラボで使用する Tensor Processing Unit(TPU)ハードウェア アクセラレータは非常に高速です。課題は、TPU を十分に活用できる速度でデータをフィードすることです。Google Cloud Storage(GCS)は非常に高いスループットを維持できますが、すべてのクラウド ストレージ システムと同様に、接続の開始時にネットワークの往復コストが発生します。そのため、データを数千もの個別のファイルとして保存することは理想的ではありません。これらのファイルをより少ない数のファイルにバッチ処理し、tf.data.Dataset の機能を使用して複数のファイルから並行して読み取ります。
Read-through
画像ファイルを読み込み、共通のサイズに変更してから 16 個の TFRecord ファイルに保存するコードは、次のノートブックにあります。ざっとお読みください。この後の Codelab では、適切な TFRecord 形式のデータが提供されるため、実行する必要はありません。
Flower pictures to TFRecords.ipynb
最適な GCS スループットを実現する理想的なデータ レイアウト
TFRecord ファイル形式
Tensorflow でデータを保存する際に推奨されるファイル形式は、protobuf ベースの TFRecord 形式です。他のシリアル化形式も使用できますが、次のように記述することで、TFRecord ファイルからデータセットを直接読み込むことができます。
filenames = tf.io.gfile.glob(FILENAME_PATTERN)
dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.map(...) # do the TFRecord decoding here - see below
最適なパフォーマンスを得るには、次のより複雑なコードを使用して、複数の TFRecord ファイルから一度に読み取ることをおすすめします。このコードは、N 個のファイルから並列で読み取り、読み取り速度を優先してデータの順序を無視します。
AUTOTUNE = tf.data.AUTOTUNE
ignore_order = tf.data.Options()
ignore_order.experimental_deterministic = False
filenames = tf.io.gfile.glob(FILENAME_PATTERN)
dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTOTUNE)
dataset = dataset.with_options(ignore_order)
dataset = dataset.map(...) # do the TFRecord decoding here - see below
TFRecord クイック リファレンス
TFRecord には、バイト文字列(バイトのリスト)、64 ビットの整数、32 ビットの浮動小数点数の 3 種類のデータを格納できます。常にリストとして保存されます。単一のデータ要素はサイズ 1 のリストになります。次のヘルパー関数を使用して、データを TFRecord に保存できます。
バイト文字列の書き込み
# warning, the input is a list of byte strings, which are themselves lists of bytes
def _bytestring_feature(list_of_bytestrings):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=list_of_bytestrings))
整数の書き込み
def _int_feature(list_of_ints): # int64
return tf.train.Feature(int64_list=tf.train.Int64List(value=list_of_ints))
浮動小数点数の書き込み
def _float_feature(list_of_floats): # float32
return tf.train.Feature(float_list=tf.train.FloatList(value=list_of_floats))
上記のヘルパーを使用して TFRecord を書き込む
# input data in my_img_bytes, my_class, my_height, my_width, my_floats
with tf.python_io.TFRecordWriter(filename) as out_file:
feature = {
"image": _bytestring_feature([my_img_bytes]), # one image in the list
"class": _int_feature([my_class]), # one class in the list
"size": _int_feature([my_height, my_width]), # fixed length (2) list of ints
"float_data": _float_feature(my_floats) # variable length list of floats
}
tf_record = tf.train.Example(features=tf.train.Features(feature=feature))
out_file.write(tf_record.SerializeToString())
TFRecord からデータを読み取るには、まず保存したレコードのレイアウトを宣言する必要があります。宣言では、名前付きフィールドに固定長リストまたは可変長リストとしてアクセスできます。
TFRecord からの読み取り
def read_tfrecord(data):
features = {
# tf.string = byte string (not text string)
"image": tf.io.FixedLenFeature([], tf.string), # shape [] means scalar, here, a single byte string
"class": tf.io.FixedLenFeature([], tf.int64), # shape [] means scalar, i.e. a single item
"size": tf.io.FixedLenFeature([2], tf.int64), # two integers
"float_data": tf.io.VarLenFeature(tf.float32) # a variable number of floats
}
# decode the TFRecord
tf_record = tf.io.parse_single_example(data, features)
# FixedLenFeature fields are now ready to use
sz = tf_record['size']
# Typical code for decoding compressed images
image = tf.io.decode_jpeg(tf_record['image'], channels=3)
# VarLenFeature fields require additional sparse.to_dense decoding
float_data = tf.sparse.to_dense(tf_record['float_data'])
return image, sz, float_data
# decoding a tf.data.TFRecordDataset
dataset = dataset.map(read_tfrecord)
# now a dataset of triplets (image, sz, float_data)
便利なコード スニペット:
単一のデータ要素の読み取り
tf.io.FixedLenFeature([], tf.string) # for one byte string
tf.io.FixedLenFeature([], tf.int64) # for one int
tf.io.FixedLenFeature([], tf.float32) # for one float
要素の固定サイズ リストの読み取り
tf.io.FixedLenFeature([N], tf.string) # list of N byte strings
tf.io.FixedLenFeature([N], tf.int64) # list of N ints
tf.io.FixedLenFeature([N], tf.float32) # list of N floats
可変数のデータアイテムを読み取る
tf.io.VarLenFeature(tf.string) # list of byte strings
tf.io.VarLenFeature(tf.int64) # list of ints
tf.io.VarLenFeature(tf.float32) # list of floats
VarLenFeature はスパース ベクトルを返します。TFRecord のデコード後に、追加のステップが必要です。
dense_data = tf.sparse.to_dense(tf_record['my_var_len_feature'])
TFRecords にオプションのフィールドを含めることもできます。フィールドの読み取り時にデフォルト値を指定すると、フィールドが存在しない場合にエラーではなくデフォルト値が返されます。
tf.io.FixedLenFeature([], tf.int64, default_value=0) # this field is optional
学習した内容
- 🤔 GCS からの高速アクセス用にデータファイルをシャーディングする
- 😓 TFRecords の書き込み方法。(もう構文を忘れてしまいましたか?(このページをチートシートとしてブックマークしておくと便利です)
- 🤔 TFRecordDataset を使用して TFRecord から Dataset を読み込む
このチェックリストを頭の中で確認してください。
6. 完了
これで、TPU にデータをフィードできるようになりました。次のラボに進んでください
- [このラボ] TPU スピードのデータ パイプライン(tf.data.Dataset と TFRecords)
- 転移学習を使用した最初の Keras モデル
- Keras と TPU を使用した畳み込みニューラル ネットワーク
- Keras と TPU での最新の convnets、squeezenet、Xception
TPU の実践
TPU と GPU は Cloud AI Platform で使用できます。
- Deep Learning VM の場合
- AI Platform Notebooks で
- AI Platform Training ジョブの場合
最後に、皆様のフィードバックをお待ちしています。このラボで不具合が見つかった場合や、改善すべき点があると思われる場合は、お知らせください。フィードバックは GitHub の問題 [フィードバック リンク] から送信できます。

|
|

