AI & 機械学習

Google Cloud 上のデータへのアクセスを迅速かつ簡単にする TensorFlow Enterprise

TensorFlowEnterprise.jpg

※この投稿は米国時間 2019 年 10 月 31 日に、Google Cloud blog に投稿されたものの抄訳です。

AI イニシアチブにおいて、データはその中心となります。簡単に言うと、ディープ ラーニング モデルをトレーニングするには、大量のデータを収集して保存する必要があります。また、GPU Cloud TPU などのアクセラレータが進歩し、その可用性が増すにつれ、保存場所からデータを取得してトレーニング プロセスに取り込むまでの速度がますます重要になってきています。

TensorFlow は、最もよく利用されている機械学習フレームワークの 1 つで、2015 年に Google によってオープンソース化されました。TensorFlow はあらゆるユーザーに対応しますが、特に Google Cloud にデプロイした場合は、TensorFlow の作成者によるエンタープライズ クラスのサポートとパフォーマンスのメリットを享受できます。これが最近、Google Cloud での AI 対応ビジネスを対象とした TensorFlow Enterprise の提供を開始した理由です。

この投稿では、TensorFlow Enterprise によってもたらされるメリットを、Google Cloud に保存されたデータへのアクセスの面から見ていきます。トレーニング データの保存に Cloud Storage を使用している場合は、「Cloud Storage リーダーの改善」セクションに進み、TensorFlow Enterprise によって Cloud Storage からのデータ スループットが 2 倍になる仕組みをご覧ください。データの保存に BigQuery を使用している場合は、「BigQuery リーダー」セクションに進み、TensorFlow Enterprise によって、TensorFlow から BigQuery データに直接アクセスして高スループットを実現できるようになる仕組みをご覧ください。

Cloud Storage リーダーの改善

TensorFlow Enterprise では、Cloud Storage から TensorFlow Dataset がデータを読み取る方法に改善が加えられています。この改善の効果を測るために、同じ TensorFlow コードを TensorFlow 1.14 と TensorFlow Enterprise で実行して、Cloud Storage から読み取られる 1 秒あたりの平均サンプル数を比較してみます。このコードは、単に tfrecord ファイルを読み取り、1 秒あたりのサンプル数を出力します。

  import tensorflow.compat.v1 as tf
import time

tf.disable_eager_execution()

def input_fn(data_files_pattern,
             batch_size,
             num_iterations=1):
  filenames = tf.io.gfile.glob(data_files_pattern)
  dataset = tf.data.Dataset.from_tensor_slices(filenames).repeat()
  dataset = dataset.apply(tf.data.experimental.parallel_interleave(
      map_func=tf.data.TFRecordDataset,
      cycle_length=10,
      sloppy=True))
  dataset = dataset.batch(batch_size, drop_remainder=False)
  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
  return dataset


def run_benchmark(data_files_pattern, num_iterations):
  batch_size = 2048
  dataset = input_fn(
        data_files_pattern=data_files_pattern,
        batch_size=batch_size)
  itr = tf.data.make_one_shot_iterator(dataset)
  size = tf.shape(itr.get_next())[0]
  with tf.Session() as sess:
    size_callable = sess.make_callable(size)
    n = 0
    mini_batch = 100
    for i in range(num_iterations // mini_batch):
      local_start = time.time()
      start_n = n
      for j in range(mini_batch):
        n += size_callable()
      local_end = time.time()
      examples_per_second = (mini_batch * batch_size) / (local_end - local_start)
      print('Processed %d entries in %f seconds. [%f] examples/s' % (
          n - start_n, local_end - local_start, examples_per_second
          ))

run_benchmark('gs://cloud-samples-data/ai-platform/fake_imagenet/train*', 20000)

出力は次のようになります。

  Processed 204800 entries in 21.885156 seconds. [9357.941099] examples/s
Processed 204800 entries in 17.571123 seconds. [11655.487321] examples/s
Processed 204800 entries in 16.429514 seconds. [12465.371641] examples/s
...

このテストで使用するデータセットは架空の ImageNet データで、 gs://cloud-samples-data/ai-platform/fake_imagenet から取得できます。

最初に TensorFlow 1.14 を搭載した Compute Engine VM、次に TensorFlow Enterprise を搭載した Compute Engine VM で同じコードを実行し、1 秒あたりの平均サンプル数を比較します。このテストでは、CPU 8 個、メモリ 64 GB の VM でコードを実行して、VM と同じリージョン内にあるリージョン Cloud Storage バケットから読み取りました。

テストの結果、大きな改善が認められました。

TensorFlow_Speed.0998018019640329.max-700x700.png

Cloud Storage からのデータ読み取り速度は、TensorFlow で提供されるさまざまなパラメータを調整することによってさらに改善できます。たとえば、前述の parallel_interleave 呼び出しを次のコードに置き換えます。

  dataset = dataset.apply(tf.data.experimental.parallel_interleave(
      map_func=tf.data.TFRecordDataset,
      cycle_length=20,
      block_length=1,
      buffer_output_elements=50000,
      prefetch_input_elements=40,
      sloppy=True))

次のようなさらなる改善を確認できます。

TensorFlow_Speed_Improvement.max-800x800.png

Cloud Storage からの読み取り速度に影響を与える要素はほかにもあります。よくある間違いが、少数の大きな tfrecord ファイルではなく、多数の小さい tfrecord ファイルを Cloud Storage に保存することです。TensorFlow による Cloud Storage からのデータ読み取りのスループットを高くするには、各ファイルが 150 MB を超えるようにデータをグループ化する必要があります。

BigQuery リーダー

TensorFlow Enterprise には BigQuery リーダーが導入されているので、BigQuery からデータを直接読み取ることができます。たとえば、次のようなコードを実行できます。

  import tensorflow as tf
from tensorflow_io.bigquery import BigQueryClient

import time

PROJECT_ID = "<YOUR-GCP-PROJECT-ID>"
DATASET_GCP_PROJECT_ID = "bigquery-public-data"
DATASET_ID = "samples"
TABLE_ID = "wikipedia"

def run_benchmark(num_iterations):
  batch_size = 2048
  client = BigQueryClient()
  read_session = client.read_session(
      "projects/" + PROJECT_ID,
      DATASET_GCP_PROJECT_ID, TABLE_ID, DATASET_ID,
      ["title",
       "id",
       "num_characters",
       "language",
       "timestamp",
       "wp_namespace",
       "contributor_username"],
      [tf.string,
       tf.int64,
       tf.int64,
       tf.string,
       tf.int64,
       tf.int64,
       tf.string],
      requested_streams=10
  )

  dataset = read_session.parallel_read_rows(sloppy=True).batch(batch_size)
  itr = dataset.make_one_shot_iterator()

  n = 0
  mini_batch = 100
  for i in range(num_iterations // mini_batch):
    local_start = time.time()
    start_n = n
    for j in range(mini_batch):
      n += batch_size
      batch = itr.get_next()

    local_end = time.time()
    print('Processed %d entries in %f seconds. [%f] examples/s' % (
        n - start_n, local_end - local_start,
        (mini_batch * batch_size) / (local_end - local_start)))

run_benchmark(10000)

BigQuery リーダーでは、データ スループットを高くするための並列データアクセス向けに、BigQuery の Storage API が使用されます。出力は次のようになります。

  Processed 204800 entries in 4.917719 seconds. [41645.322712] examples/s
Processed 204800 entries in 4.066770 seconds. [50359.375164] examples/s
Processed 204800 entries in 4.859606 seconds. [42143.334014] examples/s
...

前述のコードでは、サンプルの各バッチは Python OrderedDict として構造化されます。そのキーは read_session 呼び出しに指定された列名、値はモデルで使用される TensorFlow テンソルです。たとえば、次のようになります。

  OrderedDict([..., 
('id', <tf.Tensor: id=34, shape=(2048,), dtype=int64, numpy=array([504869,   551777, 21024884, ...])>),
...,
('title', <tf.Tensor: id=38, shape=(2048,), dtype=string, numpy=
array(['Touch typing', 'Air ioniser',
       'List of number-one songs on American Top 40 of 2009', ...], dtype=object)>),
...])

まとめ

データを保存場所から取得して機械学習トレーニング プロセスに取り込むまでの速度は、ディープ ラーニング モデル作成者の生産性にとって、ますます重要になってきています。TensorFlow Enterprise を利用することで、パフォーマンスが最適化され、データソースへのアクセスが簡単になります。また、TensorFlow Enterprise では、ユーザーが GCP での TensorFlow ワークロードの効率性を向上させるための改善が継続してサポートされます。  

詳細は、TensorFlow Enterprise についての Google のサマリーブログをご覧ください。ご利用を開始するには、TensorFlow Enterprise にアクセスしてください。