chainer: 独自datasetを定義する方法
chainerで独自データセットクラスを作るための方法を明示的に示したドキュメントが 見当たらなかったので、備忘録をかねて書く。実はとっても簡単。
- データセットにするクラスは
chainer.dataset.DatasetMixin
を継承する - 内部に持っているデータの数を返却する
__len__(self)
メソッドを実装する。このメソッドは整数でデータ数を返却する - i番目のデータを取得する
get_example(self, i)
メソッドを実装する。このメソッドは、- 画像配列などのデータ
- ラベル
の2つを返却する(return image_array, label みたいな感じで)
本当に必要なことはこのたった3つです。
Datasetのクラスを定義するタイミングで画像を全部読み込んでもよいですが、 get_example
を呼び出すタイミングで実際の読み込みを
行うのでも構いません。
実例:
import sys import random import numpy as np from PIL import Image import csv import chainer from chainer import datasets class ImageDataset(chainer.dataset.DatasetMixin): def __init__(self, normalize=True, flatten=True, train=True, max_size=200): self._normalize = normalize self._flatten = flatten self._train = train self._max_size = max_size pairs = [] with open('data/filename_label_list.tsv', newline='') as f: tsv = csv.reader(f, delimiter='\t') for row in tsv: if 'jpg' in row[0]: pairs.append(row) self._pairs = pairs def __len__(self): return len(self._pairs) def get_image(self, filename): image = Image.open('data/' + filename) new_w = self._max_size + 1 new_h = self._max_size image = image.resize((new_w, new_h), Image.BICUBIC) image_array = np.asarray(image) return image_array # type cast image_array = image_array.astype('float32') label = np.int32(label) return image_array, label def get_example(self, i): filename = self._pairs[i][0] image_array = self.get_image(filename) if self._normalize: image_array = image_array / 255 if self._flatten: image_array = image_array.flatten() else: if image_array.ndim == 2: mage_array = image_array[np.newaxis,:] image_array = image_array.astype('float32') image_array = image_array.transpose(2, 0, 1) # order of rgb / h / w label = np.int32(self._pairs[i][1]) return image_array, label
__init__
で、ファイル名とラベルのリストを読み込んでいます。ここでself._pairs
にリストの各行を入れていますが、画像データはまだ読み込んでいません__len__
はself._pairs
の項目数を返却するだけget_example
で実際の画像読み込み=> 配列化とlabelの返却を行っています(画像読み込みは、エポック数が多い場合は__init__
内で先に全部読み込んでおいたほうが早い場合もあるかもしれません)__len__
が返却する値がdatasetのサイズとみなされます。その結果、- len / minibatch_sizeがepoch内で学習されるminibatchの個数となる
- Iteratorは0番目からlen番目までの要素をget_example(i)で取得するようになる
また、注意しておいたほうが良いことが数点だけあります。
- ラベルを整数で返却する場合、
np.int32( label )
という感じでnp.int32
にキャストして返却すること. 普通のintでも回せますが、GPUを使わず、CPUのみで実行しようとするとき、labelがint型だとCUDA environment is not correctly set up
という関連の分かり辛いエラーで怒られてしまいます - 画像を普通に読み込むと、0 ~ 255 の整数データになるため、0.0 ~ 1.0に正規化すること(私は'float32'型を指定しています)
- 画像を返却する際は(色次元数, h, w) という順番に軸変換を行っておくこと。普通にPIL等でイメージを読み込むと、( h, w, 色次元数 ) という順になるため、
image_array = image_array.transpose(2, 0, 1)
などで変換が必要です - datasetクラスをinitializeする際に全てのデータを読み込んだり生成したりするのはGPUメモリ容量的に得策ではありません。initialize時にはデータのリストだけ作り、get_example内でデータの実体を読み込む/生成するようにしたほうが良いと思います。(特にデータ数が多い場合)
いろんなデータセットで楽しむきっかけになれば幸いです。