Apache SparkによるTensorFlow、PyTorchへのデータ変換の簡素化

[How to Simplify Data Conversion for Deep Learning with Apache Spark](https://databricks.com/blog/2020/06/16/simplify-data-conversion-from-apache-spark-to-tensorflow-and-pytorch.html)の翻訳です。

Petastormは、Apache Parquetフォーマットのデータセットからディープラーニングモデルに対して、シングルマシン、あるいは分散でのトレーニング、評価を実現するために、Uberから提供されている人気のあるオープンソースのライブラリです。Petastorm 0.9.0による、Apache SparkデータフレームからTensorFlowデータセット、PyTorchデータローダーへのデータ変換のサポートを発表できることを嬉しく思います。新たなSparkデータセットコンバーターAPIは、複数のデータソースからの大量データを用いた分散モデルトレーニング、推論を容易なものにします。SparkデータセットコンバーターAPIは、Yevgeni Litvin、Travis Addair (Uber)と、Xiangrui Meng、Weichen Xu、Liang Zhang (Databricks)のコラボレーションによって開発されました。

# なぜディープラーニングにおけるデータ変換が難しいのか？

あらゆるディープラーニングパイプラインにおいてキーとなるステップは、データをDLフレームワークの入力フォーマットに変換する部分です。Apache Sparkは最も人気のあるビッグデータフレームワークです。Apache Sparkデータフレームからディープラーニングフレームワークにデータを変換することは退屈なものです。例えば、特徴量のカラム、ラベルのカラムを保持するApache SparkデータフレームをTensorFlowデータセットファイルフォーマットに変換するためには、ユーザーはApache Sparkデータフレームを分散ファイルシステム上にParquetフォーマットで保存し、Petastormのようなサードパーティツールを用いて変換したデータをロードするか、spark-tensorflow-connectorを用いて直接TFRecordファイルとして保存し、TFRecordDatasetを用いてロードし直す必要があります。いずれのアプローチも、中間データファイルを管理するために20行以上のコードが必要となり、異なるパーシングの文法に依存し、Sparkデータフレームでカラムを取り扱うためには追加の注意が必要となります。これらのエンジニアリングにおける軋轢は、データサイエンティストの生産性を阻害します。

# ソリューション概要

Databricksはこれらの面倒なデータ変換プロセスステップをシンプルなものにするために、Petastorに対して新たな[SparkデータセットコンバーターAPI](https://petastorm.readthedocs.io/en/latest/readme_include.html#spark-dataset-converter-api)を寄贈しました。新たなAPIによって、Sparkデータフレームから[TensorFlowデータセット](https://www.tensorflow.org/datasets)や[PyTorchデータローダー](https://pytorch.org/docs/stable/data.html)(デフォルトのパラメーター)への変換が数行のコードで済みます。

```py:Python
from petastorm.spark import SparkDatasetConverter, make_spark_converter

# Specify the cache directory
spark.conf.set(SparkDatasetConverter.PARENT_CACHE_DIR_URL_CONF, 'file:///dbfs/tmp/…')

df = spark.read...

converter = make_spark_converter(df)    # create the converter

with converter.make_tf_dataset() as dataset:    # convert to TensorFlow Dataset
        # Training or inference code with dataset 
        ...
with converter.make_torch_dataloader() as dataloader:    # convert to PyTorch DataLoader
        # Training or inference code with dataloader 
        ...
```

# Sparkデータセットコンバーターは何をするのか？

SparkデータセットコンバーターAPIは以下の機能を提供します。

- **キャッシュ管理** コンバーターはSparkデータフレームを分散ファイルシステムにキャッシュし、ベストエフォートでインタプリタがexitしたときにキャッシュファイルを削除します。明示的な削除APIも提供されています。
- **出力データセットをカスタマイズするための豊富なパラメーター** ユーザーはbatch_size、workers_countを設定することで出力データセットをカスタマイズ、コントロールすることができ、ベストなI/Oパフォーマンスを達成するために事前フェッチを行います。
- **pandasデータフレームに対して定義される変換関数** 多くのディープラーニングデータセットには、Sparkデータフレームのバイナリー列としてロードできる画像、音声、動画が含まれています。これらのバイナリー列はディープラーニングモデルに入力する前にデコードが必要となります。コンバーターは、デコーディングロジックを指定するための変換関数に対するフックを提供します。変換関数は入力として、Sparkデータフレームから変換されたpandasデータフレームを受け取り、デコードされたデータを含むpandasデータフレームを返却する必要があります。
- **MLlibベクトルのハンドリング** プリミティブなデータ型に加え、コンバーターはSparkデータフレームをキャッシュする前にSpark MLlibベクトル型を自動で配列カラムに変換することで、Spark MLlibのベクトル型をサポートします。変換関数の中で一次元の配列を多次元配列にリシェイプすることもできます。
- **リモートデータロード** コンバーターをSparkのワーカーノードにpickleすることができ、ワーカーノード上でTensorFlowデータセットあるいはPyTorchデータローダーを作成するために使用することができます。パラメーターで、特定のシャードあるいは全てのデータセットを読み込むかを指定することができます。
- **シングルノードから分散コンピューティングへの移行が容易**　お使いのシングルノードの推論コードを分散推論に移行する際、データハンドリングのコードの変更は不要であり、普通にSparkで動作します。分散トレーニングにおいては、シャードのインデックスとシャードの総数をAPIのパラメーターに追加するだけで済みます。我々のエンドツーエンドのサンプルノートブックでは、どのようにシングルノードのコードを、[Horovod](https://github.com/horovod/horovod)を用いて、どのように分散推論、分散トレーニングに移行するのかを説明しています。

詳細はリソースセクションのリンクをチェックしてみてください。

# 使い始めてみる

Databricks Machine Learningランタイム7.0以降をも強いて、以下でリンクされているエンドツーエンドのサンプルノートブックを試してみてください。

**AWSノートブック**

- [SparkからTensorFlowへのデータ変換の簡素化](https://github.com/taka-yayoi/public_repo/tree/main/petastorm_aws_tensorflow)
- [Simplify data conversion from Spark to PyTorch](https://databricks.com/notebooks/simple-aws/petastorm-spark-converter-pytorch.html)

**Azureノートブック**

- [Simplify data conversion from Spark to TensorFlow](https://databricks.com/notebooks/simple-azure/petastorm-spark-converter-tensorflow.html)
- [Simplify data conversion from Spark to PyTorch](https://databricks.com/notebooks/simple-azure/petastorm-spark-converter-pytorch.html)

# 謝辞

この機能を実現するための議論、レビューをしてくれたPetastormの作者であるUberのYevgeni LitvinとTravis Addairに感謝いたします！

# リソース

- Databricksのドキュメントおよびエンドツーエンドのサンプル([AWS](https://docs.databricks.com/applications/deep-learning/data-prep/petastorm.html)|[Azure](https://docs.microsoft.com/en-us/azure/databricks/applications/deep-learning/data-prep/petastorm))
- [Petastorm GitHub Homepage](https://github.com/uber/petastorm)
- [Petastorm SparkDatasetConverter API documentation](https://petastorm.readthedocs.io/en/latest/api.html#module-petastorm.spark.spark_dataset_converter)

### Databricks 無料トライアル

[Databricks 無料トライアル](https://databricks.com/jp/try-databricks)