はじめに

昨年の12月に発表されて話題をさらったpix2pixモデルだが、既にgithub上に幾つかの実装コードが存在する。
しかしCPUのみで動くようなコードが無いので作ってみた。

環境

OS:Ubuntu14.04
GPU:GTX1070（今回は使わない）
CUDA:8.0 RC（今回は使わない）
cuDNN:5.1（今回は使わない）
python：2.7.6
chainer：1.20.0.1
など

参考文献等

まずpix2pixの原論文は以下。
https://arxiv.org/abs/1611.07004
①今回参考にしたコードは２つ。まずmattyaさんのDCGANsコード。
https://github.com/mattya/chainer-DCGAN
②次にpreferdのpix2pixコード。
https://github.com/pfnet-research/chainer-pix2pix
こちらの画像取得部分のコード（facade_dataset.py）はほぼそのまま使わせてもらってます。

原論文に対する変更点

原論文に対する主な変更点は以下。

generator及びdiscriminatorともに層の数、パラメータ等を大幅に削減
入出力画像のサイズは28x28に圧縮
L2ノルムなどの正則化は使わない
dropoutも使わない
（3,4に関してはbatch normalizationを使っているので、ある程度代替できてるか？）

モデルの概略

Generatorは以下のように変更した。

Discriminatorは以下のように変更した。

コード

コードは以下

train_pix2pix_cpu.py

import numpy as np
from PIL import Image
import os
import pylab

import chainer
from chainer import optimizers
from chainer import serializers
from chainer import Variable
import chainer.functions as F
import chainer.links as L

image_dir = './images_plane_s' #input data file
out_image_dir = './out_images_miniplane' #output image file
out_model_dir = './out_models_miniplane' #output model file

nz2 = 10 #dimention of noise
batchsize = 100
n_epoch = 200

# read all images
fs = os.listdir(image_dir)
print(len(fs))
dataset = []
for fn in fs:
    f = open('%s/%s' % (image_dir, fn), 'rb')
    img_bin = f.read()
    dataset.append(img_bin)
    f.close()
print(len(dataset))

# function for output image
def clip_img(x):
    return np.float32(-1 if x < -1 else (1 if x > 1 else x))
zvis = (np.random.uniform(-1, 1, (100, nz2)).astype(np.float32))
try:
    os.mkdir(out_image_dir)
    os.mkdir(out_model_dir)
except:
    pass


class Generator(chainer.Chain):
    def __init__(self):
        super(Generator, self).__init__(
            l0s=L.Linear(nz2, 7 * 7 * 64),
            dc1s=L.Deconvolution2D(64, 32, 4, stride=2, pad=1),
            dc2s=L.Deconvolution2D(32, 3, 4, stride=2, pad=1),
            bn0s=L.BatchNormalization(7 * 7 * 64),
            bn1s=L.BatchNormalization(32),
        )

    def __call__(self, z, test=False):
        h = self.l0s(z)
        h = self.bn0s(h, test=test)
        h = F.relu(h)
        h = F.reshape(h, (z.data.shape[0], 64, 7, 7))
        h = self.dc1s(h)
        h = self.bn1s(h, test=test)
        h = F.relu(h)
        x = self.dc2s(h)
        return x


class Discriminator(chainer.Chain):
    def __init__(self):
        super(Discriminator, self).__init__(
            c0s=L.Convolution2D(3, 32, 4, stride=2, pad=1),
            c1s=L.Convolution2D(32, 64, 4, stride=2, pad=1),
            l2s=L.Linear(7 * 7 * 64, 2),
            bn0s=L.BatchNormalization(32),
            bn1s=L.BatchNormalization(64),
        )

    def __call__(self, x, test=False):
        h  = self.c0s(x)
        h = F.leaky_relu(h, slope=0.2)
        h = self.c1s(h)
        h = self.bn1s(h, test=test)
        h = F.leaky_relu(h, slope=0.2)
        l = self.l2s(h)
        return l


gen = Generator()
dis = Discriminator()

optimizer_gen = optimizers.Adam(alpha=0.0002, beta1=0.5)
optimizer_dis = optimizers.Adam(alpha=0.0002, beta1=0.5)
optimizer_gen.setup(gen)
optimizer_dis.setup(dis)


# training loop
for epoch in range(0, n_epoch):
    sum_loss_dis = np.float32(0)
    sum_loss_gen = np.float32(0)

    for i in range(0, len(dataset), batchsize):

        x2 = np.zeros((batchsize, 3, 28, 28), dtype=np.float32)
        for j in range(batchsize):
            # try:
            rnd = np.random.randint(len(dataset))
            img = np.asarray(Image.open(image_dir + '/'+ fs[rnd]).convert('RGB')).astype(np.float32)
            img = img.transpose(2, 0, 1)
            x2[j, :, :, :] = (img[:, :, :] - 128.0) / 128.0

        z = Variable(np.random.uniform(-1, 1, (batchsize, nz2)).astype(np.float32))
        x = gen(z)
        yl = dis(x)
        d_0 = Variable(np.zeros(batchsize, dtype=np.int32))
        d_1 = Variable(np.ones(batchsize, dtype=np.int32))
        loss_gen = F.softmax_cross_entropy(yl, d_0)
        loss_dis = F.softmax_cross_entropy(yl, d_1)

        yl2 = dis(x2)
        t2 = Variable(np.zeros(batchsize, dtype=np.int32))
        loss_dis += F.softmax_cross_entropy(yl2, t2)

        optimizer_gen.zero_grads()
        loss_gen.backward()
        optimizer_gen.update()

        optimizer_dis.zero_grads()
        loss_dis.backward()
        optimizer_dis.update()

        sum_loss_gen += loss_gen.data
        sum_loss_dis += loss_dis.data

    print('epoch end', epoch, sum_loss_gen / len(dataset), sum_loss_dis / len(dataset))

    if epoch % 10 == 0:
        #draw image
        pylab.rcParams['figure.figsize'] = (16.0, 16.0)
        pylab.clf()
        vissize = 100
        z = zvis
        z[50:, :] = (np.random.uniform(-1, 1, (50, nz2)).astype(np.float32))
        z = Variable(z)
        x = gen(z, test=True)
        x = x.data
        for i_ in range(100):
            tmp = ((np.vectorize(clip_img)(x[i_, :, :, :]) + 1) / 2).transpose(1, 2, 0)
            pylab.subplot(10, 10, i_ + 1)
            pylab.imshow(tmp)
            pylab.axis('off')
        pylab.savefig('%s/vis_%d.png' % (out_image_dir, epoch))

        #save model
        serializers.save_hdf5("%s/dcgan_modeloss_dis_%d.h5" % (out_model_dir, epoch), dis)
        serializers.save_hdf5("%s/dcgan_modeloss_gen_%d.h5" % (out_model_dir, epoch), gen)
        serializers.save_hdf5("%s/dcgan_state_dis_%d.h5" % (out_model_dir, epoch), optimizer_dis)
        serializers.save_hdf5("%s/dcgan_state_gen_%d.h5" % (out_model_dir, epoch), optimizer_gen)

また画像の読み込み等は以下のコード。これは上記②のコードを一部変更している。

facade_dataset2.py

import os

import numpy
from PIL import Image
import six

import numpy as np

from io import BytesIO
import os
import pickle
import json
import numpy as np

import skimage.io as io

from chainer.dataset import dataset_mixin

# download `BASE` dataset from http://cmp.felk.cvut.cz/~tylecr1/facade/
class FacadeDataset(dataset_mixin.DatasetMixin):
    def __init__(self, dataDir='./facade/base', data_range=(1,300)):
        print("load dataset start")
        print("    from: %s"%dataDir)
        print("    range: [%d, %d)"%(data_range[0], data_range[1]))
        self.dataDir = dataDir
        self.dataset = []
        for i in range(data_range[0],data_range[1]):
            #print("file_name", dataDir+"/cmp_b%04d.jpg"%i)
            img = Image.open(dataDir+"/cmp_b%04d.jpg"%i)
            label = Image.open(dataDir+"/cmp_b%04d.png"%i)
            w,h = img.size
            #print("w , h = ", w, h)
            #r = 286/min(w,h)
            r = 31 / float(min(w, h))
            # resize images so that min(w, h) == 286
            #print("r = ", r)
            #print("min(w,h) = ", min(w,h))
            img = img.resize((int(r*w), int(r*h)), Image.BILINEAR)
            #print("img.size = ", img.size)
            label = label.resize((int(r*w), int(r*h)), Image.NEAREST)
            #debug
            img = np.asarray(img)
            #print("type(img) = ", type(img))

            #print("img.shape = ", img.shape)
            
            img = np.asarray(img).astype("f").transpose(2,0,1)/128.0-1.0
            label_ = np.asarray(label)-1  # [0, 12)
            label = np.zeros((12, img.shape[1], img.shape[2])).astype("i")
            for j in range(12):
                label[j,:] = label_==j
            #print("label.shape = ", label.shape)
            #print("label[0][[10][10] = ", label[0][10][10])
            self.dataset.append((img,label))
        print("load dataset done")
    
    def __len__(self):
        return len(self.dataset)

    # return (label, img)
    def get_example(self, i, crop_width=28):
        _,h,w = self.dataset[i][0].shape
        x_l = np.random.randint(0,w-crop_width)
        x_r = x_l+crop_width
        y_l = np.random.randint(0,h-crop_width)
        y_r = y_l+crop_width
        # print("np.max(self.dataset[i][1][:,y_l:y_r,x_l:x_r]) = ", np.max(self.dataset[i][1][:,y_l:y_r,x_l:x_r]))
        # print("np.max(self.dataset[i][0][:,y_l:y_r,x_l:x_r]) = ", np.max(self.dataset[i][0][:,y_l:y_r,x_l:x_r]))
        # print("np.min(self.dataset[i][1][:,y_l:y_r,x_l:x_r]) = ", np.min(self.dataset[i][1][:,y_l:y_r,x_l:x_r]))
        # print("np.min(self.dataset[i][0][:,y_l:y_r,x_l:x_r]) = ", np.min(self.dataset[i][0][:,y_l:y_r,x_l:x_r]))

        return self.dataset[i][1][:,y_l:y_r,x_l:x_r], self.dataset[i][0][:,y_l:y_r,x_l:x_r]

現在諸事情で個人用GitHubが使えなくなっているが、復活したらそちらにもupします。

結果

学習前の出力画像と500回学習した場合の出力画像は以下。

学習前（左からイラスト、本物の画像、生成された画像）

500回学習後（左からイラスト、本物の画像、生成された画像）