TorchData
pytorchでデーラローダーを作成する方法は色々ありますが,最近(2022年前半)にTorchDataがアナウンスされました.
でもあまり情報が出回っておらず,検索してもcsvを読み込む程度の説明しか見当たりません.データパイプの分岐や処理などをどうやるのか,何に使えるのかがよく分からないので,記事を書いてみました.
リンク先の記事ではtorchdataの簡単な説明から始めて,物体検出用のデータセットに適用するということをしています.画像は画像フォルダに,bbox情報はcsvに,というアノテーションデータをtorchdataのdatapipeを使って読み込みます.データ拡張にはalbumentationsを使います.
WebDataset
torchdata登場以前からWebDatasetというものがあるというのを聞いており,こちらもpytorchのデータ読み込みのための便利なライブラリです.
こちらもあまり情報が出回っておらず,検索しても使い方をすぐに学べるような情報がありません.画像認識のための画像ファイルとラベル情報のファイルをセットにしたような例しかなく,そもそもshardってなにという概念からしてよく分からないので,記事を書いてみました.
リンク先の記事ではwebdatasetのいろいろなshard作成方法を説明しています.セマンティックセグメンテーション用の実画像とラベル画像のペアを読み込んでshardを作る例もあります.またそれらのshardを読み込んでwebdatasetのパイプラインで読み込み・デコード・データ拡張などを行う方法も説明しています.データ拡張にはalbumentationsを使います.
で,torchdataとwebdatasetの違いは?
そもそものpytorchのdatasetオブジェクトは,いろいろ使いにくいところがあります.それを補うのがこれらのライブラリですが,ポリシーが全く違います.そのため,適材適所で使っても良いですし,併用もできるでしょう.
torchdataの使いどころ
いろいろなデータセットをダウンロードして使う場合,展開したディレクトリ構造やアノテーションファイルのフォーマットは個別のデータセットに依存しています.画像はサブディレクトリに保存されており,アノテーション情報はcsvやjsonに,というようにバラバラの場合,これまでは自分でコードを書く必要がありました.torchdataはそれらを統一に扱えるAPIです.
処理は,あらかじめ各処理をパイプラインでつなげてdatapipeというオブジェクトを構築するという方法をとっています.このパイプライン処理を,途中で分岐させたり,分岐させたパイプラインをまたつなげたり,ということもできます.またもともとデータが(画像ファイルとjsonというような)別々の場合には,それぞれを処理するパイプラインを作成し,最後につなげて,それからDataLoaderでサンプルを取得する,ということができます.
また以下のリンク先ではさまざまなvision datasetに対してどう使うのかのコードがあります.
これを眺めて見ると,S3やgoogle driveやHTTPなどweb上のtarやzipを直接指定して,それを展開しながらdatapipeに通してDataLoaderでサンプルを取得する,という一連の流れが書いてあります.tarやzipをいちいち展開する必要がないというのは魅力です(が,巨大なtarファイルをいちいちダウンロードする必要はないなので,使い方次第ですが).
webdatasetの使いどころ
webdatasetは,shardと呼ばれる複数のtarファイルを作成し,それらからデータを読み込むことに特化しています.ですので,まずshardを作成しなければなりません.いちいちshardにする必要がなければ,webdatasetは必要ないでしょう.
する必要がある場合というのが,大量の画像ファイルを読み込むような場合です.たとえば数千万枚の画像を学習する必要がある場合,容量の関係からローカルのHDDにしか保存出来ないでしょう(SSDには入り切らないとする).この場合に画像をランダムに選択して学習すると,HDDへのランダムアクセスが発生し,非常に効率が悪くなります(つまり遅い).個別のファイルをオープン・クローズするコストも馬鹿になりません(多分...).また数TB程度のHDDではinodeを使い切ってしまい,容量は余っているのにデータセットのファイルをそれ以上保存できない,というようなことも発生します.
その場合,数千万枚の画像ファイルを数十個のtarファイルに分けて保存しておき,これらのtarファイルの中からdataloaderが読み込むようにするのがwebdatasetです.こうするとHDDからの読み込みとしてはシーケンシャルリードになり(ランダムアクセスよりも速い),ファイル数は少なくinodeの問題はなくなり,ファイルのオープン・クローズの回数も非常に少なくなります.
読み込むshardはローカルHDDに限りません.数百TBや数PBのshardはローカルHDDに保存できないので,ローカルネットワークのNASやweb上のS3などから取得することになります.この場合,tarファイルを読み込みながら(ダウンロードしながら)DataLoaderでサンプルを取得することができます.(そのため「ペタバイト級も扱える」というような説明を見かけることになります).
処理は,あらかじめ各処理をパイプラインでつなげてオブジェクトを構築するという方法をとっています.これはtorchdataと同じですが,違う点はパイプライン処理を途中で分岐させたりつなげたり,ということは想定していない点です.shardを作成する時点で,1つのサンプルの情報を(同じキーを持ち拡張子だけが異なる複数のファイルとしてtarの中に)まとめておく必要があります.
torchdataとwebdataset
torchdataをwebdatasetのshard作成に利用
torhdataとwebdatasetは設計ポリシーが異なるので,両方を同時に使う必要はありません.併用する一つの案としては,webdatasetのshard作成のときにtorchdataを利用してデータをまとめる,という方法があります.
webdatasetのshardをtorchdataで読み込み
shardにまとめるというのはwebdatasetのコンセプトですが,そのtarファイルをtorchdataで読み込むこともできます.torchdataはwebdatasetフォーマットのtarファイルを読み込む関数があります.