AI & 機械学習

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

2019年11月15日

https://storage.googleapis.com/gweb-cloudblog-publish/images/TensorFlowEnterprise.max-2600x2600.jpg

Google Cloud Japan Team

※この投稿は米国時間 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)

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

読み込んでいます...

このテストで使用するデータセットは架空の 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 バケットから読み取りました。

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

https://storage.googleapis.com/gweb-cloudblog-publish/images/TensorFlow_Speed.0998018019640329.max-700x70.max-700x700.png

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

読み込んでいます...

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

https://storage.googleapis.com/gweb-cloudblog-publish/images/TensorFlow_Speed_Improvement.max-800x800.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 が使用されます。出力は次のようになります。

読み込んでいます...

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

読み込んでいます...

まとめ

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

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

投稿先