1. 概要
TPU は非常に高速です。トレーニング データのストリームも、トレーニングのスピードについていく必要があります。このラボでは、tf.data.Dataset API を使用して GCS からデータを読み込み、TPU にフィードする方法を学びます。
このラボは、「TPU での Keras」シリーズのパート 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] を選択します。この Codelab では、ハードウェア アクセラレーションによるトレーニングをサポートする強力な TPU(Tensor Processing Unit)を使用します。ランタイムへの接続は、初回実行時に自動的に行われます。右上の [接続] ボタンを使用することもできます。
ノートブックの実行
セルをクリックして Shift+Enter キーを押すことで、セルを 1 つずつ実行します。[ランタイム] > [すべて実行] を使用して、ノートブック全体を実行することもできます。
目次
すべてのノートブックには目次があります。左側の黒い矢印を使用して開くことができます。
非表示のセル
一部のセルはタイトルのみが表示されます。これは 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=16, 000 個のコアが必要になりますが、通常は不可能です。最大の GPU には約 4,000 個のコアがあります。一方、TPU は、MXU のコンピューティング単位に、最低限のハードウェア(bfloat16 x bfloat16 => float32
乗アキュムレータのみ)を使用します。これらは非常に小さいため、TPU は 128x128 MXU で 16K を実装し、この行列の乗算を一度に処理できます。
図: MXU シストリック アレイ。コンピューティング要素は乗算累積器です。1 つの行列の値が配列に読み込まれます(赤いドット)。他の行列の値は配列(グレーのドット)を介して流れます。縦線は値を上方に伝播します。水平線は部分和を伝播します。データが配列を流れ、右側から行列乗算の結果が得られることを、ユーザーが確認してください。
それに加えて、ドット積が MXU で計算されている間、中間合計は単に隣接する計算ユニット間で流れます。メモリやレジスタ ファイルとの間で保存や取得を行う必要はありません。その結果、TPU のサイクロル アレイ アーキテクチャは、行列乗算の計算において、GPU と比較して密度と電力の点で大きな利点があり、速度の面でも無視できない利点があります。
Cloud TPU
Google Cloud Platform で 1 つの「Cloud TPU v2」をリクエストすると、PCI に接続された TPU ボードを備えた仮想マシン(VM)が提供されます。TPU ボードには、4 つのデュアルコア TPU チップが搭載されています。各 TPU コアは、VPU(ベクトル プロセッシング ユニット)と 128x128 MXU(MatriX 乗算ユニット)を備えています。この「Cloud TPU」は、通常はネットワークを介してリクエスト元の VM に接続されます。全体像は次のようになります。
図: ネットワークに接続された「Cloud TPU」アクセラレータを備えた VM。「Cloud TPU」自体は、4 つのデュアルコア TPU チップを搭載した PCI 接続 TPU ボードを備えた VM で構成されています。
TPU Pod
Google のデータセンターでは、TPU がハイ パフォーマンス コンピューティング(HPC)相互接続に接続されているため、1 つの非常に大規模なアクセラレータのように見えます。Google ではこれを Pod と呼び、最大 512 個の TPU v2 コアまたは 2,048 個の TPU v3 コアを含めることができます。
イラスト: TPU v3 Pod。HPC Interconnect で接続された TPU ボードとラック。
トレーニング中には、all-reduce アルゴリズムを使用して TPU コア間で勾配が交換されます(all-reduce についてはこちらの説明をご覧ください)。トレーニング中のモデルは、大きなバッチサイズでトレーニングすることでハードウェアを活用できます。
イラスト: Google TPU の 2D トロイダル メッシュ 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(Accelerated Linear Algebra コンパイラ)というコンパイラが、コンピューティング ノードの TensorFlow グラフを TPU マシンコードに変換します。このコンパイラは、コードとメモリ レイアウトに対して多くの高度な最適化も行います。コンパイルは、処理が TPU に送信されると自動的に実行されます。XLA をビルドチェーンに明示的に含める必要はありません。
例: TPU で実行するには、TensorFlow プログラムで定義された計算グラフがまず XLA(高速線形代数コンパイラ)表現に変換され、次に XLA によって TPU マシンコードにコンパイルされます。
Keras での TPU の使用
TPU は、Tensorflow 2.1 以降の Keras API を通じてサポートされています。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」ユーティリティで作成されたディープ ラーニング VM)で、パラメータなしで機能します。これらのシステムは、TPU_NAME 環境変数により TPU の場所を認識します。TPU を手動で作成する場合は、使用している VM で TPU_NAME 環境変数を設定するか、明示的なパラメータTPUClusterResolver(tp_uname, zone, project)
を指定してTPUClusterResolver
を呼び出します。TPUStrategy
は、分布と「all-reduce」勾配同期アルゴリズムを実装する部分です。- 戦略はスコープを通じて適用されます。モデルは strategy scope() 内で定義する必要があります。
tpu_model.fit
関数は、TPU トレーニングの入力として tf.data.Dataset オブジェクトを想定しています。
TPU 移行の一般的なタスク
- TensorFlow モデルにデータを読み込む方法は数多くありますが、TPU では
tf.data.Dataset
API を使用する必要があります。 - TPU は非常に高速であり、実行時にデータの取り込みがボトルネックになることがよくあります。TPU パフォーマンス ガイドには、データのボトルネックを検出するためのツールや、パフォーマンスに関するその他のヒントが記載されています。
- int8 または int16 の数値は int32 として扱われます。TPU には、32 ビット未満で動作する整数ハードウェアはありません。
- 一部の Tensorflow オペレーションはサポートされていません。リストはこちらです。幸いなことに、この制限はトレーニング コード、つまりモデルのフォワード パスとバックワード パスにのみ適用されます。データ入力パイプラインでは、CPU で実行されるため、すべての TensorFlow オペレーションを使用できます。
tf.py_func
は TPU でサポートされていません。
4. データの読み込み
今回は花の写真のデータセットを使用します。目標は、花を 5 種類に分類する方法を学ぶことです。データの読み込みは tf.data.Dataset
API を使用して実行されます。まず、API について説明します。
ハンズオン
次のノートブックを開き、セルを実行(Shift+Enter)し、「作業が必要」のラベルがある場合は、手順に沿って操作してください。
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 の機能を使用して複数のファイルから並列に読み取ります。
リードスルー
画像ファイルを読み込んで一般的なサイズに変更し、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 クイック リファレンス
TFRecords には、バイト文字列(バイトのリスト)、64 ビットの整数、32 ビットの浮動小数点数の 3 種類のデータが保存できます。それらは常にリストとして格納され、1 つのデータ要素はサイズ 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())
TFRecords からデータを読み取るには、まず保存したレコードのレイアウトを宣言する必要があります。宣言では、固定長のリストまたは可変長のリストとして名前付きフィールドにアクセスできます。
TFRecords からの読み取り
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 からすばやくアクセスする
- 😝? TFRecord の書き方について確認しました。(構文を忘れた場合は、ブックマークして、このページをクイック リファレンスとして使用してください)。
- 🤔 TFRecordDataset を使用して TFRecords から Dataset を読み込む
このチェックリストの内容を頭に入れておいてください。
6. 完了
これで、TPU にデータをフィードできるようになりました。次のラボに進んでください。
- [このラボ] TPU スピードのデータ パイプライン: tf.data.Dataset と TFRecords
- 転移学習を使用した最初の Keras モデル
- Keras と TPU を使用した畳み込みニューラル ネットワーク
- Keras と TPU を使用した最新の convnets、squeezenet、Xception
TPU の実践
TPU と GPU は Cloud AI Platform で使用できます。
ぜひフィードバックをお寄せください。このラボで問題が見つかった場合や、改善が必要だと思われる場合は、お知らせください。フィードバックは、GitHub の問題 [フィードバック リンク] からお送りいただけます。
|