Pytorch で DCGAN をつくったよ

DCGAN を二年ぶりに実装しました。

いつも MNIST も面白くないので、 MNIST にアルファベットが加わった EMNIST をデータセットとして用いました。

www.nist.gov

f:id:dette:20180502060329p:plain — Emnist の画像例

こんな感じの画像が10万オーダーで格納されています。大きさは Mnist と同じ (28, 28) の grayscale の画像です。

DCGAN とは

DCGAN はニューラルネットワークの生成モデルである GAN (Generative Adversarial Networks) の一種です。 DCGANでは画像を生成するモデル (generator) だけではなく、画像が偽物(generatorが生成したもの）かデータセット上にある本物の画像かを判定するモデル (dicriminator) も同時に訓練するという枠組みのモデルです。こうすることで生成モデルはより判別モデルを騙すような高度な画像を生成し、それに応じて判別モデルはより精緻に偽画像を判定するようになり、競い合いながら学習が進んでいきます。

アルゴリズムに関しては以前の記事とかぶるので DCGANをChainerで実装 - nykergoto’s blog や、とても良くまとまっているこちらなど参考にしていただけると良いかなと思います。

じっそう

pytorh での実装について。まずはモデル部分から。 generator / dicsriminator ともに似たような感じなので generator だけ

class Generator(nn.Module):
    """
    ランダムベクトルから shape = (1, 28, 28) の画像を生成するモジュール
    Note: はじめ Linear な部分には Batchnorm1d を入れていなかったが学習が Discriminator に全く追いつかず崩壊した。
    Generator のネットワーク構成はシビアっぽい
    """

    def __init__(self, z_dim=100):
        super().__init__()
        self.z_dim = z_dim

        self.flatten_dim = 5 ** 2 * 256
        self.fc = nn.Sequential(*[
            nn.Linear(z_dim, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU(True),
            nn.Linear(1024, self.flatten_dim),
            nn.BatchNorm1d(self.flatten_dim),
            nn.ReLU(True)
        ])

        self.upsample = nn.Sequential(
            make_layer(256, 128, kernel=3, stride=1, padding=0, deconv=True), # 5 -> 7
            make_layer(128, 64, kernel=4, stride=2, deconv=True), # 7 -> 14
            nn.ConvTranspose2d(64, 1, kernel_size=4, stride=2, padding=1) # 14 -> 28
        )

    def forward(self, input_tensor):
        h = self.fc(input_tensor)
        h = h.view(-1, 256, 5, 5)
        x = self.upsample(h)
        x = F.sigmoid(x)
        return x

工夫があるとすれば make_layer という形で nn.Module を返すヘルパ関数を用意した、という点でしょうか。

def make_layer(in_channels, out_channels, kernel=2, stride=2, padding=1, use_batchnorm=True, deconv=False):
    layers = []

    if deconv:
        conv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=kernel, stride=stride, padding=padding)
    else:
        conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel, stride=stride, padding=padding)

    layers.append(conv)
    if use_batchnorm:
        bn = nn.BatchNorm2d(out_channels)
        layers.append(bn)
    activate = nn.ReLU(True)
    layers.append(activate)

    return nn.Sequential(*layers)

今回は3層程度なので直接記述しても大したことは無いですが、 resnet50 のように大きなネットワークになるとこのようにある塊をひとつの関数で作るようにすることは必要になるのかなーと思っています。（このやり方で良いのかは微妙ですが……みんなどうしているのだろうか…）

つぎに学習を実行する Solver クラスを作ってみました。

class Solver(object):
    """
    abstract class for solver
    """

    def __init__(self, models, device=None):
        """
        :param tuple[torch.nn.Module] models: 学習中に重みが更新されるモデルの配列
        :param str device: train 時のデバイス. `"cuda"` or `"cpu"`
        """
        self.models = models

        if isinstance(device, str):
            if device == "cuda" and torch.cuda.is_available():
                self.device = torch.device("cuda")
            else:
                self.device = torch.device("cpu")
        else:
            self.device = device

        self.callbacks = None

    def start_epoch(self, epoch):
        for model in self.models:
            model.to(self.device)
            model.train()

        for c in self.callbacks:
            c.initialize(self.models, self.device)
            c.start_epoch(epoch)

    def end_epoch(self, epoch, logs):
        for c in self.callbacks:
            c.end_epoch(epoch, logs)

    def fit(self, train_iterator, valid_iterator=None, epochs=10, callbacks=None, initial_lr=0.01,
            optimizer_params=None):
        """
        訓練データを用いてモデルの fitting を行う
        :param train_iterator: Iterable Object. 各 Iteration で (x_train, label_train) の tuple を返す
        :param valid_iterator:
        :param int epochs:
        :param list[Callback] callbacks:
        :param float initial_lr:
        :param dict | None optimizer_params:
        :return:
        """
        if optimizer_params is None:
            optimizer_params = {}
        self.set_optimizers(lr=initial_lr, params=optimizer_params)
        self.callbacks = callbacks

        for epoch in range(1, epochs + 1):
            self.start_epoch(epoch)
            logs = self.train(train_iterator)
            self.end_epoch(epoch, logs)
        return self

    def set_optimizers(self, lr, params):
        raise NotImplementedError()

    def train(self, train_iterator):
        logs = {}
        count = 0
        total = len(train_iterator)
        for batch_idx, (x, _) in tqdm(enumerate(train_iterator), total=total):
            count += 1
            logs = self._forward_backward_core(x, _, logs)

        for k, v in logs.items():
            if "loss" in k:
                logs[k] /= count
        return logs

    def _forward_backward_core(self, x, t, logs):
        raise NotImplementedError()

このクラスを作ったのにはちょっと理由があって、 pytorch のチュートリアルでは大概のものが epoch を引数にとって学習を行うクラス def train(epoch) のみを実行して model 部分などは global スコープにある変数を参照する形式をとっていてとても気持ちが悪かった、というのがあります。（なんとなく関数の引数以外の変数を使うのってこわくないですか）

で前使っていた keras を参考に training ループを回す部分とエポックの始まり終わりの定期実行を実装したものを書いてみた、というのが先の Solver クラスです。

最適化モデルの設定や forward-backward の定義はモデルや学習スキームによって変化するのでその部分は継承クラスで定義、という形式をとっています。

で、実装してみて思ったことなのですが……
結局こうやって枠にはめてしまうと pytorch の持っている柔軟なネットワークの構成ができる利点が薄れてしまってあまり良くないなあという印象を持ちました。

そもそも、こういう型にはめて記述できるものは keras などもっと便利に隠匿してくれるパッケージで書くことができますしそちらのほうが楽なんだろうなーということです。 pytorch を使うような場合にはそれこそ def train(epoch), def test(epoch) のように関数ベースで定義してループもごりごり書くほうが結果としてわかりやすくなるのかもしれません。この辺はまだ pytorch 初心者なのでもっと良い書き方があるのだろうなとは思っているのでいろんなコードを読んで勉強していきたいなという所存です。