AI Labの大田（@ciela）です。近頃はリサーチエンジニアとして研究成果を最大化するためのエンジニアリングに責任を持って日々活動しています。昨年もPyTorchに関する記事を書かせていただきましたが、今回もPyTorch、特にPyTorchでの学習データセットの取り扱いについてのお話になります。これまで部署内のプロダクトと関わる上でTensorFlowに浮気することも何度かありましたが、PyTorchに対する感覚を鈍らせたくはなく、今後もPyTorch関連の話題には頑張ってキャッチアップしていく所存です。

深層学習データセット取扱の課題

昨今の深層学習に用いられるデータセット内のデータ数は一般的に大規模です。実際に学習する上ではデータセットをシャッフルしつつ繰り返しロードし、場合によっては適切な前処理やData Augmentationだってしなくてはなりません。それらの処理を並列化することで、学習にかかる時間を少しでも減らそうと涙ぐましい努力をされている方も多いことでしょう。その一方で、時には数万件にも及ぶファイルに対してランダムアクセスしつつ何度もロードするアクセスパターンは入出力効率を低下させてしまうためストレージデバイスにとって優しくありません。この矛盾を解決するために、格納プロトコルによって規格化された独自アーカイブファイルにデータセットを変換し、それを一定サイズに分割（シャーディング）した形で保管するアプローチが採用されたりします。ある程度まとまった容量のファイルをシーケンシャルにロードすることで、ストレージの入出力効率を高めることが狙いです。

深層学習フレームワークの代表格であるTensorFlowにおいてはプロトコルバッファ形式でデータをシリアライズ・アーカイブするTFRecordという仕組みが提供されています。例えばGoogle Cloud Platform (GCP)のCloud Storageに200〜300MB程度の容量ごとにシャーディングされた学習データのTFRecordファイルを配置しておけば、あとはそのパスパターンを指定するだけで学習データをシーケンシャルにロードすることができ、ファイル単体ごとにランダムにロードする場合と比べて学習速度の向上が期待できます。学習データの前処理パイプライン構築に最適なApache Beam（GCP Dataflow）からも簡単に書き出せるため、GCP上でTensorFlowを利用している方にとっては既に当たり前の仕組みかもしれません。一方で、TensorFlowと肩を並べる深層学習フレームワークであるPyTorchではどういったアプローチが考えられているのでしょうか。昨年8月、PyTorch公式ブログにEfficient PyTorch I/O library for Large Datasets, Many Files, Many GPUsという記事が投稿されており、記事内ではWebDatasetというTensorFlowでのTFRecordに相当する仕組みが紹介されています。本記事ではこのWebDatasetについて簡単にまとめ、実際のデータセットでの利用例をお見せしようと思います。

WebDataset

WebDatasetは国際会議BigData 2019にて採択されたHigh Performance I/O For Large Scale Deep Learningという論文で定義されている大規模な深層学習データセットを取り扱うための機構およびその実装です。PyTorchのIterableDatasetクラスを継承実装しており、データベース・リモートサーバ・リアルタイムログからストリーミングでデータを取得するような場合など、データ全体に対するランダムアクセスが困難・不可能でバッチサイズがロード済みデータに依存してしまうような環境下でのデータセットを取り扱うことに適した仕組みになっています。TensorFlowのTFRecord同様、WebサーバやクラウドストレージにWebDatasetとしてデータをシャーディングして配置しておくことでその状態のまま学習ルーチンへ取り込むことが可能となります。2021年3月現在はPyTorchとは別リポジトリで開発が進められているものの、既にRFCとしてプロポーザルが上がっており将来的にはPyTorch内部に取り入れられるかもしれません。

WebDatasetの大きな特徴として、実体がPOSIX tar形式でアーカイブされていることが挙げられます。*nixなコンピュータを利用する方にとっては馴染み深すぎる形式ですね。データセット内の各データとラベルを同一のbasenameを持ったファイルとして対応付けた状態でアーカイブしており、tarコマンドで中身を確認したり展開したりといったことも容易に行なえます。読み込み用の特別な仕組みを必要としないため、ポータブルなデータセットアーカイブとしての役割も果たしていると言えるでしょう。下記は手元で作成したWebDatasetの中身を閲覧してみた例です。

WebDatasetを読み込んだあとの、シャッフル・torch.Tensor化・バッチ化などの学習中前処理はTensorFlow tf.data.Datasetと似たような書き方で定義できます。具体的にはメソッドチェーンを利用して処理を繋げていくfluent interfaceとなっており、下記のような形で記述できます。

この他にも、メディアファイルの拡張子に基づく自動デコード・エンコード、取得したリソース（シャード、レコード共に）のローカルキャッシュ、などアーカイブされたデータセットを取り扱う上で嬉しい工夫が凝らされてます。詳しくはGitHubリポジトリのJupyter Notebook集によくまとまっていますのでこちらをご参照ください。

既存のデータセットをWebDatasetへ変換

実際に既存のデータセットをWebDatasetへ変換してみたいと思います。例としてBehance Artistic Media（BAM）データセットを利用してみましょう。BAMデータセットの実体はSQLiteファイルとなっており、DBレコードのsrcカラムに指定されたBehance CDNのURLを参照することで画像をロードできますが、学習時にいちいちインターネットアクセスをしていては学習ルーチンの効率が悪すぎます。さらに既にアクセスが不可能となってしまっている（404）画像に対する例外処理が必要であったりと、そのままでは扱いが若干ややこしいデータセットです。画像を予めダウンロードしてローカルストレージに保存しておき、各DBレコードのURLをファイルパスへ変換しつつランダムアクセスするのがPyTorchでの素直な学習ワークロードになるかと思われますが、ここではそれをWebDatasetに置き換えてみましょう。

今回は下記のようなSQLを発行して得られる、BAMデータセットDBのcrowd_labelsテーブルによる属性ラベルを画像と対応付けたレコードをデータセットとして考え、これをWebDataset化したいと思います。BAMデータセットの主要なDBスキーマについては以前の記事で簡単に紹介しておりますのでよろしければご覧ください。

クエリ

レコードサンプル

今回は少量データで試すために各属性ごとのデータは10件ずつに絞っています。さらにこのデータセットをWebDatasetに変換するために下記のようなコードを実行してみましょう。

URLにアクセスして取得できたバイト列をリソースの拡張子に合わせてWebDatasetに一件一件格納していきます。さらにシャーディングを試すために圧縮前の最大ファイルサイズはあえて10MBに制限しています（デフォルト3GB）。実行した結果、手元の環境に4つのtarballファイルが生成されました。データセットのWebDataset化自体はこれで完了です。簡単ですね。

あとは読み込みです。上記のローカルファイルをファイルパスとして直接読み込んでも良いのですが、せっかくなのでこれらが外部サーバに配置された状態を想定し、HTTP経由でウェブリソースとして取得するシミュレーションを行ってみましょう。Python3のhttp.serverモジュールを利用して簡易ウェブサーバを立上げWebDatasetをローカル配信してみます。

この状態で下記のコードを実行してみましょう。WebDatasetに格納されているデータセットから、画像はRGB3チャンネルでサイズ224×224のtorch.Tensor、ラベルはstrとしてそれぞれデコード・前処理を行って読み込むコードになります。

実行結果は下記のようになりました。無事WebDatasetを読み込むことに成功し、画像とその属性ラベルがバッチ化された状態で取得できたことが確認できました。

おわりに

本記事では深層学習データセットを取り巻く課題とWebDatasetによるアプローチを簡単に解説しました。PyTorchで大規模なメディアデータセットを取り扱う上での指針となるのではないでしょうか。

最近ではTensorFlow Datasets（TFDS）というライブラリも出てきており、既存のデータセットをより簡単にTFRecordとして読み書きする仕組みが整っていたりします。こちらは名前の通りTensorFlow向けの仕組みですが、同様の思想のもとPyTorch向けにWebDatasetを書き出すような何らかのシステムを整えることは可能なはずです。独自データセットをTraining-readyな状態で利用したいのにTensorFlow向けにしか提供されていない…というような危険な宗教戦争の火種を潰せますね。

AI Labではこのような研究のためのエンジニアリングをより加速させていきたいと思います。今後は採用も考えていければと思うのでこういった活動に興味のある方はどうぞご連絡ください。