ChainerのTrainerを使ってFineTuning

python chainer 機械学習 人工知能

どうもこんにちは

chainerなんか早くなったらしいですね。

とても素晴らしい技術力だと思いますが、個人的には高速化よりも詳しい使い方が書いたドキュメントや日本語ドキュメントの充実に力を入れてほしい気がします。
（GPU128枚も持ってないし…）
128枚もGPUを使って何かするようなところって大きな企業かでかい研究機関しかないと思うし、そういったところが使うかと言うと使わない気もする。 PyTorchとかchainerに似ているフレームワークも出てきていて正直chainerどうなるんだろうという感じはしていますが、PaintsChainerとかあったり一応日本では流行ってるみたいなので使ってみました。

Trainer

Trainerってなんだって話ですが、今まではバッチ処理を自分で書いたりしなければならなかったのがTrainerによって抽象化されて書く必要がなくなりました。 いまどれだけの進捗でどれくらいのlossになっててAccuracyはどれくらいだと言うことも割りと簡単にできるようになったみたいなので今回はこのTrainerを使ってみたいと思います。 学習させるところまでできたので記録としてかいてあります。 間違ってるところとか、こういうふうに書いたほうがかっこいいぜみたいなのあったら教えていただけるとありがたいです。

使い方

上のサンプルコードを参考に書いていきます。詳しいことはこのサンプルとドキュメントを見るとわかりやすいかと思います。 まず、事前にやっておくこととして、FineTuningをしたいのでモデルの重みを取ってきておきます。 modelはこちらにたくさんあるので好きなのを選びましょう。

caffeのモデルですが、chainerではcaffeのモデルをロードできるような仕組みがCaffefunctionと言うかたちで実装されています。 caffefunctionの詳しい使い方は今回書きませんが、caffeのモデルを読み込むのはすごく時間がかかるので、一度読み込んだらPickleなどでdumpしておきましょう。 重みの読み込みは以下ののコードの場合load_model.pyということに書いてあります。重みの読み込みはこちらで書いていることと同じことを行っています。

手順としては、
重みを読み込む→トレーニングセットと、バリデーションセットに分ける→optimizer,updaterの設定をする→trainerでトレーニング

と言った感じです。 Trainerで学習させるときにGPUを使う場合はモデルにたいしてto_gpu()を行う必要があります。ここで結構ハマりました。

import chainer
import pickle
import chainer.links as L
from chainer import training
from chainer.training import extensions
from preprocess import ImagePreprocess
import vgg16
from load_model import load_caffe_model
import os
import pandas as pd


def trainer(csv_file, root, batch_size, lorderjob, image_size, gpu):

    # load model
    model = vgg16.Vgg16()
    model_weight = pickle.load(open('../data/models/vgg16model', 'rb'))
    load_caffe_model(model_weight, model)
    model = L.Classifier(model)
    print('finish loading model')

    if gpu >= 0:
        chainer.cuda.get_device(0).use()
        model.to_gpu()

    # create train set and validation set
    csv_data = pd.read_csv('../data/train.tsv', delimiter='\t')
    dataset = []
    for r in csv_data.iterrows():
        dataset.append((r[1][0], r[1][1]))
    train, val = chainer.datasets.split_dataset_random(dataset, 9000)
    train_data = ImagePreprocess(train, root, image_size)
    val_data = ImagePreprocess(val, root, image_size, False)

    train_iter = chainer.iterators.MultiprocessIterator(
        train_data, batch_size, n_processes=lorderjob
    )
    val_iter = chainer.iterators.MultiprocessIterator(
        val_data, batch_size, repeat=False, n_processes=lorderjob
    )

    # setup optimizer
    optimizer = chainer.optimizers.SGD()
    optimizer.setup(model)

    # setup updater and trainer
    updater = training.StandardUpdater(train_iter, optimizer, device=0)
    trainer = training.Trainer(updater, (10, 'epoch'), out='result')

    trainer.extend(extensions.Evaluator(val_iter, model, device=0))
    trainer.extend(extensions.dump_graph('main/loss'))
    trainer.extend(extensions.snapshot(), trigger=(5, 'epoch'))
    trainer.extend(extensions.snapshot_object(
        optimizer, 'optimizer_iter_{.updater.iteration}'), trigger=(5, 'epoch'))
    trainer.extend(extensions.LogReport())
    trainer.extend(extensions.PrintReport(
        ['epoch', 'main/loss', 'validation/main/loss',
         'main/accuracy', 'validation/main/accuracy']
    ), trigger=(1, 'epoch'))
    trainer.extend(extensions.ProgressBar(update_interval=10))
    trainer.run()

    print('Training is finished!')

DataMixinの拡張をして、画像をミニバッチごとに読み込む

import chainer
import random
import cv2


class ImagePreprocess(chainer.dataset.DatasetMixin):

    def __init__(self, dataset, root, img_size, random=True):
        self.base = chainer.datasets.LabeledImageDataset(dataset, root)
        self.img_size = img_size
        self.random = random

    def __len__(self):
        return len(self.base)

    def get_example(self, i):
        img_size = self.img_size
        image, label = self.base[i]
        image = image.transpose(1, 2, 0)
        image = cv2.resize(image, (img_size, img_size))
        image = image.transpose(2, 0, 1)
        _, h, w = image.shape

        if self.random:
            # Randomly imga region and flip the image
            if random.randint(0, 1):
                image = image[:, :, ::-1]

        image *= (1.0 / 255.0)  # Scale to [0, 1]
        return image, label

preprocess.pyというファイルにDatasetMixinクラスを拡張したものを作っています。このクラスを拡張して、get_exampleをいい感じにいじるとミニバッチごとに画像を読み込んでくれるのでメモリの節約をすることができます。 DataAugmentationもこれを使うことでできます。上のコードではエポックごとに画像をランダムでフリップさせることによって、画像のデータを増やすといったことをしています。

またこのクラスのオブジェクトをMultiprocessIteratorに渡すことでCPUでミニバッチ分の画像を読み込み、GPUで学習といったことができるようになります。

以上のようにすれば学習が進んでいくかと思います。 Trainerは抽象化されているので理解するのに少し時間がかかるかと思いますが、使えるようになると便利なので使っていくといいのではないでしょうか。 特に、CPUで画像読み込みながらGPUで学習するとかいったコードを自分で書くのは大変そうなのでそのあたりもTrainer使うと便利そうです。