gotutiyan’s blog

プログラミング関連の話題を中心に

【Pytorch】ミニバッチ学習に便利なDataSet・DataLoaderの使い方

はじめに

深層学習によって学習を行う際には,ミニバッチ化して学習させることが一般的です.本記事では,pytorchで提供されているDataSetとDataLoaderという機能を用いてミニバッチ化を実現する方法について書きます.

ミニバッチ化とは

簡単にミニバッチ化について復習します.ミニバッチ化とは,訓練データをそれなりに小さいサイズで分割し,それによって得られた小規模な訓練データ(ミニバッチ)で訓練することです.分割する際には,バッチサイズが必要です.例えばバッチサイズが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])]
'''

こちらも,特に理由がない限りはTrueにしておくと良いと思います.

なぜDataLoaderなのか

ミニバッチは,言ってしまえば[data[i*batch_size : (i+1)*batch_size] for i in range(len(data) / batch_size]みたいに,リストのスライスを使うだけでも実現できます.でもなぜDataLoaderを使うのかというと,個人的には以下の点が良いと思っています.

  1. shuffleやdrop_lastの引数が便利
  2. 分割後のデータをtensorで返してくれる
  3. ミニバッチ化してることが一目で分かる

特に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を用いてミニバッチ化する方法について詳しく説明しました.

参考になれば幸いです.