PyTorchではDataLoaderを使うことで読み込んだデータから自動でミニバッチを作成することができます。 DataLoaderを使いこなすことで、ニューラルネットワークの学習部分を簡単に書くことができます。 本記事ではPyTorchのDataLoaderがミニバッチを作成する仕組みについて解説します。

目次

DataLoaderでは以下のような引数を受け付けます。

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
           batch_sampler=None, num_workers=0, collate_fn=None,
           pin_memory=False, drop_last=False, timeout=0,
           worker_init_fn=None)

本記事では主にdataset、sampler、batch_sampler、collate_fnが関連します。

Dataset

datasetはmap-styleとiterable-styleの二種類あり、これらは利用目的に合わせて使い分けます。

Map-style dataset

__getitem__と__len__を実装しているランダムアクセス可能なオブジェクトです。 __getitem__の引数が整数ではなくても良いのでmap-styleと呼んでいます。

Iterable-style dataset

__iter__を実装し、IterableDatasetを継承したシーケンシャルアクセスが可能なオブジェクトです。 map-styleのようにランダムアクセスが難しい場合（例えば、データがメモリに乗り切らない場合）に適しています。 iterable-styleではユーザがデータの順序を実装する必要があります。

Dataloderによるデータのミニバッチ化

PyTorchのDataLoaderは読み込んだデータを自動でミニバッチ化する機能が有効となっています。 これPyTorchのドキュメントではAutomatic batchingと呼ばれています。 この章ではAutomatic batchingが有効時と無効時での挙動の違いについて解説します。

Automatic batching有効時の挙動

DataLoaderでは個別のデータを自動でミニバッチ化できます。 このとき、DataLoaderはそれぞれのdatasetに対して、以下のようにbatch_samplerからミニバッチを構成するデータのリストをcollate_fnへ与えます。

Map-style:

for indices in batch_sampler:
    yield collate_fn([dataset[i] for i in indices])

Iterable-style:

dataset_iter = iter(dataset)
for indices in batch_sampler:
    yield collate_fn([next(dataset_iter) for _ in indices])

デフォルトのcollate_fnは次のような挙動となります。

  • 最初の次元にbatchを表す次元を追加する。
  • NumPyのarrayやPythonの数値のlistをPyTorchのTensorに変換する。
  • データの型はそのまま保持する。例えば、入力がdict型であれば、出力はdict型のまま、ただし、valueはTensorとして返す。

Automatic batching無効時の挙動

すでにバッチ化されたデータを読み込みたい、ミニバッチ学習を利用したくない、等の場合にはautomatic batchingを無効化することもできます。 batch_sizeとbatch_samplerがともにNoneのとき (batch_samplerのデフォルト値はNone)、automatic batchingは無効化されます。 Automatic batchingが無効化されているときは、samplerはデータを一つサンプリングしてcollate_fnへ与えます。

Map-style:

for index in sampler:
    yield collate_fn(dataset[index])

Iterable-style:

for data in iter(dataset):
    yield collate_fn(data)

おわり

本記事ではPyTorchのDataLoaderにおけるミニバッチ化の仕組みについて解説しました。 automatic batchingはデフォルトで有効になっているため、DataLoaderを使いこなすことでユーザは少ない手間で学習部分を記述できます。 またautomatic batchingを無効にした時の挙動についても説明しました。


関連記事






最近の記事