【Pytorch】ミニバッチ学習に便利なDataSet・DataLoaderの使い方
はじめに
深層学習によって学習を行う際には，ミニバッチ化して学習させることが一般的です．本記事では，pytorchで提供されているDataSetとDataLoaderという機能を用いてミニバッチ化を実現する方法について書きます．
- はじめに
- ミニバッチ化とは
- DataSetとDataLoader
- DataSetの実装
- DataLoaderの実装
- DataLoaderを使う理由
- __getitem__()の返り値を辞書にする
- おわりに
ミニバッチ化とは
簡単にミニバッチ化について復習します．ミニバッチ化とは，訓練データをそれなりに小さいサイズで分割し，それによって得られた小規模な訓練データ（ミニバッチ）で訓練することです．分割する際には，バッチサイズが必要です．例えばバッチサイズが2だと，全てのデータを2個ずつの組みに分けることになります．
より具体的な例を見てみます．与えられた数字が偶数か奇数かを判定するモデルを作るとしましょう．偶数だと0，奇数だと1を出力させます．いま，訓練データが
# 入力 X = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] # 出力（教師） t = [0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
のように与えられているとします．2つのリストの同じindexに格納されているデータが，入力と出力のペアになります．例えば，X[5]=5, t[5]=1なので，入力5に対する出力は1，つまり奇数であることが分かります．また，訓練データ数は10です．ここで，バッチサイズ=2でミニバッチ化を行うと，
# 入力1 [0, 1] # 出力1 [0, 1] # 入力2 [2, 3] # 出力2 [0, 1] # 入力3 [4, 5] # 出力3 [0, 1] # 入力4 [6, 7] # 出力4 [0, 1] # 入力5 [8, 9] # 出力5 [0, 1]
のように分割されます．（ここでは，データをシャッフルするなどはせず，単純に先頭から区切っています．）このように小規模になったデータを用いて学習を回していきます．
蛇足になりますが，エポックの話をしましょう．今回の例では，データ数10に対してバッチサイズを2としたので，ミニバッチは5つ作成されます．このミニバッチを用いて学習するとき，ミニバッチ5つ分の学習を行えば，分割前のデータ数に相当する学習を行ったことになります．このように，分割前のデータ1回分の学習に相当する学習を1単位として1エポックと呼びます．
DataSetとDataLoader
さて，本題に入っていきます．
pytorchでは，DataSetとDataLoaderを用いることで，ミニバッチ化を簡単に実装できます．
DataSetは，元々のデータを全て持っていて，ある番号を指定されると，その番号の入出力のペアをただ一つ返します．クラスを使って実装します．
DataLoaderは，DataSetのインスタンスを渡すことで，ミニバッチ化した後のデータを返します．これはpytorchに元から用意されている関数を呼ぶだけなので，実装することはほぼありません．
まずはDataSetから実装します．
DataSetの実装
DataSetを実装する際には，クラスのメンバ関数として__len__()と__getitem__()を必ず作ります．
__len__()は，len()を使ったときに呼ばれる関数です．
__getitem__()は，array[i]のように[ ]を使って要素を参照するときに呼ばれる関数です．これが呼ばれる際には，必ず何かしらのindexが指定されているので，引数にindexの情報を取ります．また，入出力のペアを返すように設計します．今回はカンマ区切りで返すことにしましょう．
冒頭でミニバッチ化を説明した時と同じように，入力値が偶数か奇数かを判定するモデルを想定します．
class DataSet: def __init__(self): self.X = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] # 入力 self.t = [0, 1, 0, 1, 0, 1, 0, 1, 0, 1] # 出力 def __len__(self): return len(self.X) # データ数(10)を返す def __getitem__(self, index): # index番目の入出力ペアを返す return self.X[index], self.t[index]
さて，実際にこのDataSetがどのような振る舞いをするか試してみましょう．
dataset = DataSet() print('全データ数:',len(dataset)) # 全データ数: 10 print('3番目のデータ:',dataset[3]) # 3番目のデータ: (3, 1) print('5~6番目のデータ:',dataset[5:7]) # 5~6番目のデータ: ([5, 6], [1, 0])
DataSetクラスのインスタンスが，len()を受け付けていて，[index]による参照も行えることが確認できます．また，[index]による参照では，入出力のペアが得られていることも確認します．例えば3番目のデータは(3, 1)なので，入力が3で，出力は1，つまり奇数であることが分かります．
DataLoaderの実装
次はDataLoaderです．これはpytorchのモジュールに含まれるtorch.utils.data.DataLoader()を使います．この関数の引数には， DataSetのインスタンスと，バッチサイズをbatch_size=で渡します．（他にもパラメータは指定できますが，後で紹介することにします．）
バッチサイズは2としましょう．
# さっき作ったDataSetクラスのインスタンスを作成 dataset = DataSet() # datasetをDataLoaderの引数とすることでミニバッチを作成． dataloader = torch.utils.data.DataLoader(dataset, batch_size=2)
これだけで，分割されたデータが得られています．分割後のデータは，for文で取り出すことができます．
for data in dataloader: print(data) ''' 出力: [tensor([0, 1]), tensor([0, 1])] [tensor([2, 3]), tensor([0, 1])] [tensor([4, 5]), tensor([0, 1])] [tensor([6, 7]), tensor([0, 1])] [tensor([8, 9]), tensor([0, 1])] '''
tensor()というのは，pytorchの独自の型です．計算グラフというのを保持するので，学習時に必ず変換する型です． 一応，これで最小限のミニバッチ化ができたので，学習を回すことができます．適当に10エポック学習するとすれば
epoch = 10 model = #何かしらのモデル for _ in range(epoch): for data in dataloader: X = data[0] t = data[1] y = model(X) # lossの計算とか
などと書けるでしょう．また，
for data in dataloader: X = data[0] t = data[1]
は，素直に
for X, t in dataloader:
とも書けます．
shuffle
ここまでで得られたミニバッチの入力（X）に注目すると，先頭から順に0,1,2,3...となっています．この原因は，分割時に単に先頭から切り取っているからです．しかし，実際にはランダムに選んでくるのが望ましいです．そのようなとき，shuffle=Trueとすることで実現できます．
dataset = DataSet() dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True) for data in dataloader: print(data) ''' 出力: [tensor([4, 1]), tensor([0, 1])] [tensor([0, 7]), tensor([0, 1])] [tensor([9, 3]), tensor([1, 1])] [tensor([6, 5]), tensor([0, 1])] [tensor([8, 2]), tensor([0, 0])] '''
順番がバラバラになりました．さらに，同じインスタンスであっても，取り出すたびに順番が変化します．
dataset = DataSet() dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True) print('1回目') for data in dataloader: print(data) print() print('2回目') for data in dataloader: print(data) print() print('3回目') for data in dataloader: print(data) ''' 出力: 1回目 [tensor([0, 3]), tensor([0, 1])] [tensor([4, 7]), tensor([0, 1])] [tensor([1, 9]), tensor([1, 1])] [tensor([5, 8]), tensor([1, 0])] [tensor([2, 6]), tensor([0, 0])] 2回目 [tensor([5, 9]), tensor([1, 1])] [tensor([0, 2]), tensor([0, 0])] [tensor([4, 1]), tensor([0, 1])] [tensor([6, 8]), tensor([0, 0])] [tensor([3, 7]), tensor([1, 1])] 3回目 [tensor([4, 2]), tensor([0, 0])] [tensor([6, 0]), tensor([0, 0])] [tensor([3, 1]), tensor([1, 1])] [tensor([7, 8]), tensor([1, 0])] [tensor([5, 9]), tensor([1, 1])] '''
特に理由がなければ，Trueにしておくべき引数です．
drop_last
ここまでは，データ数10に対してバッチサイズが2でした．10は2で割り切れるので，綺麗に分割できました．ですが，バッチサイズを3にすればどうでしょうか．
dataset = DataSet() dataloader = torch.utils.data.DataLoader(dataset, batch_size=3) for data in dataloader: print(data) ''' 出力: [tensor([0, 1, 2]), tensor([0, 1, 0])] [tensor([3, 4, 5]), tensor([1, 0, 1])] [tensor([6, 7, 8]), tensor([0, 1, 0])] [tensor([9]), tensor([1])] '''
このように，最後に1組だけ仲間はずれになります．このようなとき，drop_last=Trueを指定することで，仲間はずれを除去できます．
dataset = DataSet() dataloader = torch.utils.data.DataLoader(dataset, batch_size=3, drop_last=True) for data in dataloader: print(data) ''' 出力: [tensor([0, 1, 2]), tensor([0, 1, 0])] [tensor([3, 4, 5]), tensor([1, 0, 1])] [tensor([6, 7, 8]), tensor([0, 1, 0])] '''
テストデータに対しては必ずFalseにすべきです．一部のデータが使われない状態で評価値を計算するのは危険です．訓練データに対してはTrueにしたほうが良いかもしれません．本来，それぞれのバッチには多様なクラスが含まれるべきです．しかし，今回の例のように偶然データが1つしかないようなミニバッチが作成されると，そのミニバッチに含まれていたクラスの情報が支配的になる可能性があります．
DataLoaderを使う理由
ミニバッチは，言ってしまえば[data[i*batch_size : (i+1)*batch_size] for i in range(len(data) / batch_size]みたいに，リストのスライスを使うだけでも実現できます．でもなぜDataLoaderを使うのかというと，個人的には以下の点が良いと思っています．
特に2.について補足ですが，分割前のデータをリストで持っていても，numpyのndarrayで持っていても，とにかくtensorにして返してくれます．例えばpandasで入力を受けてモデルに流す時などはndarrayになると思いますが，分割前のデータ型をあまり気にしないで使用できるところは大きいと思います．
ndarrayでも動作することを確かめたい人は，プログラムの先頭でimport numpy as npを書いた上で，DataSetの
def __init__(self): self.X = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] self.t = [0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
を
def __init__(self): self.X = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) self.t = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
に書き換えてみてください．
__getitem__()の返り値を辞書にする
この項目は蛇足かもしれませんが，せっかくなので書いておきます．
DataSetのところで，
def __getitem__(self, index): return self.X[index], self.t[index]
と書いていましたが，
def __getitem__(self, index): return {'X': self.X[index], 't': self.t[index]}
とすることで，dataloaderから辞書型で取り出せます．
dataset = DataSet() dataloader = torch.utils.data.DataLoader(dataset, batch_size=2) for data in dataloader: print(data) # data['X']やdata['t']で参照可能に ''' 出力: {'X': tensor([0, 1]), 't': tensor([0, 1])} {'X': tensor([2, 3]), 't': tensor([0, 1])} {'X': tensor([4, 5]), 't': tensor([0, 1])} {'X': tensor([6, 7]), 't': tensor([0, 1])} {'X': tensor([8, 9]), 't': tensor([0, 1])} '''
一つの方法として，見やすくなると思います．
おわりに
今回は，DataSetとDataLoaderを用いてミニバッチ化する方法について詳しく説明しました．
