chainer: 独自datasetを定義する方法

f:id:mizti:20170113130134p:plain

f:id:mizti:20171013202945p:plain

chainerで独自データセットクラスを作るための方法を明示的に示したドキュメントが 見当たらなかったので、備忘録をかねて書く。実はとっても簡単。

  1. データセットにするクラスは chainer.dataset.DatasetMixinを継承する
  2. 内部に持っているデータの数を返却する __len__(self) メソッドを実装する。このメソッドは整数でデータ数を返却する
  3. 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内でデータの実体を読み込む/生成するようにしたほうが良いと思います。(特にデータ数が多い場合)

いろんなデータセットで楽しむきっかけになれば幸いです。