高性能版GANの「styleGAN」で本物そっくりの画像を生成してみた【keras・機械学習】

今回は論文で紹介されてたNVIDIAが開発したstyleGANを実装してみた。

普通のGANとは生成過程も違うし、生成画像の出来の精度も比較にならないぐらい高くて、驚いた。
仕事で使う機会があったので、その生成過程をまとめてく。

目次
1.styleGANについて
2.styleGANコード詳細
3.訓練
4.生成画像
5.まとめ

1.styleGANについて

styleGANはNVIDIAが開発した、本物と見分けがつかないくらいの画像が作れる、超高精度のGAN。
qiita.com

"Progressive-Growing of GANs”というGANの亜種のgeneratorの部分を発展させたもの。

メインのメカニズムは、低い解像度の層から順に学習して、高精度の画像の生成していく仕組みらしい。

【従来のGANとstyleGANの違い】
f:id:trafalbad:20190707152347j:plain

generatorの部分がAdaptive Instance Normalization (AdaIN)でかなり改造してあるのが普通のGANと大きく異なる点。

ちなみにdiscriminatorは普通のGANと同じ。

【Generator部分】
f:id:trafalbad:20190707152336j:plain

2.styleGANコード詳細

input画像

入力画像はgoogle検索でもってきた景色の画像3枚を120枚に増幅。
f:id:trafalbad:20190707154813p:plain

Macの動画ソフト「Quick time player」で動画にして、ffmpegで静止画に変換。120枚くらいに増やした。

$ ffmpeg -i 元動画.avi -ss 144 -t 148 -r 24 -f image2 %06d.jpg

かなり簡単に画像が大量に作れるので便利。

他にもopencvで動画から、静止画を作るやり方もある。

input画像等、ハイパーパラメータの条件はこんな感じ

・画像の合計120枚

・shape=(256, 256, 3)

・Batch =10

・255で割って正規化

・1000～2000エポック

各景色画像30枚ずつで、計120枚を10batchずつ回して訓練してく。

ちなみにoptimizerを含めて、今回のパラメータは、サイズが256の時にこのパラメータでうまくいった。

けど、512とか1024は同じパラメータでは上手くいくかわからない。

Generator

generatorはinputが3つあって、2つは論文にあるように、ノイズを入れるところになってる。

AdaINが普通のGANとの違いがでかい。

from AdaIN import AdaInstanceNormalization

im_size = 256
latent_size = 512

def g_block(inp, style, noise, fil, u = True):

    b = Dense(fil)(style)
    b = Reshape([1, 1, fil])(b)
    g = Dense(fil)(style)
    g = Reshape([1, 1, fil])(g)

    n = Conv2D(filters = fil, kernel_size = 1, padding = 'same', kernel_initializer = 'he_normal')(noise)

    if u:
        out = UpSampling2D(interpolation = 'bilinear')(inp)
        out = Conv2D(filters = fil, kernel_size = 3, padding = 'same', kernel_initializer = 'he_normal')(out)
    else:
        out = Activation('linear')(inp)

    out = AdaInstanceNormalization()([out, b, g])
    out = add([out, n])
    out = LeakyReLU(0.01)(out)

    b = Dense(fil)(style)
    b = Reshape([1, 1, fil])(b)
    g = Dense(fil)(style)
    g = Reshape([1, 1, fil])(g)

    n = Conv2D(filters = fil, kernel_size = 1, padding = 'same', kernel_initializer = 'he_normal')(noise)

    out = Conv2D(filters = fil, kernel_size = 3, padding = 'same', kernel_initializer = 'he_normal')(out)
    out = AdaInstanceNormalization()([out, b, g])
    out = add([out, n])
    out = LeakyReLU(0.01)(out)

    return out

def generator():

    inp_s = Input(shape = [latent_size])
    sty = Dense(512, kernel_initializer = 'he_normal')(inp_s)
    sty = LeakyReLU(0.1)(sty)
    sty = Dense(512, kernel_initializer = 'he_normal')(sty)
    sty = LeakyReLU(0.1)(sty)

    inp_n = Input(shape = [im_size, im_size, 1])
    noi = [Activation('linear')(inp_n)]
    curr_size = im_size
    while curr_size > 4:
        curr_size = int(curr_size / 2)
        noi.append(Cropping2D(int(curr_size/2))(noi[-1]))

    inp = Input(shape = [1])
    x = Dense(4 * 4 * 512, kernel_initializer = 'he_normal')(inp)
    x = Reshape([4, 4, 512])(x)
    x = g_block(x, sty, noi[-1], 512, u=False)

    if(im_size >= 1024):
        x = g_block(x, sty, noi[7], 512) # Size / 64
    if(im_size >= 512):
        x = g_block(x, sty, noi[6], 384) # Size / 64
    if(im_size >= 256):
        x = g_block(x, sty, noi[5], 256) # Size / 32
    if(im_size >= 128):
        x = g_block(x, sty, noi[4], 192) # Size / 16
    if(im_size >= 64):
        x = g_block(x, sty, noi[3], 128) # Size / 8

    x = g_block(x, sty, noi[2], 64) # Size / 4
    x = g_block(x, sty, noi[1], 32) # Size / 2
    x = g_block(x, sty, noi[0], 16) # Size
    x = Conv2D(filters = 3, kernel_size = 1, padding = 'same', activation = 'sigmoid')(x)
    return Model(inputs = [inp_s, inp_n, inp], outputs = x)

Discriminator

Discriminatorは従来のと変わらない。

今回は、1024サイズにも対応できるようにするために、有名なGANのDiscriminatorを使った。

def discriminator():
    inp = Input(shape = [im_size, im_size, 3])

    x = d_block(inp, 16) #Size / 2
    x = d_block(x, 32) #Size / 4
    x = d_block(x, 64) #Size / 8

    if (im_size > 32):
       x = d_block(x, 128) #Size / 16

    if (im_size > 64):
        x = d_block(x, 192) #Size / 32

    if (im_size > 128):
        x = d_block(x, 256) #Size / 64

    if (im_size > 256):
        x = d_block(x, 384) #Size / 128

    if (im_size > 512):
        x = d_block(x, 512) #Size / 256

    x = Flatten()(x)
    x = Dense(128)(x)
    x = Activation('relu')(x)
    x = Dropout(0.6)(x)
    x = Dense(1)(x)
    
    return Model(inputs = inp, outputs = x)

disganとadganの実装

disganでdiscriminatorを訓練。adganでgeneratorを訓練する。

G = generator()
D = discriminator()

# Dは更新して、Gは更新しない
D.trainable = True
for layer in D.layers:
    layer.trainable = True

G.trainable = False
for layer in G.layers:
    layer.trainable = False

ri = Input(shape = [im_size, im_size, 3])
dr = D(ri)

gi = Input(shape = [latent_size])
gi2 = Input(shape = [im_size, im_size, 1])
gi3 = Input(shape = [1])
df = D(G([gi, gi2, gi3]))
da = D(ri)
disgan = Model(inputs=[ri, gi, gi2, gi3], outputs=[dr, df, da])

# Gは更新、Dは更新しない
D.trainable = False
for layer in D.layers:
    layer.trainable = False

G.trainable = True
for layer in G.layers:
    layer.trainable = True

gi = Input(shape = [latent_size])
gi2 = Input(shape = [im_size, im_size, 1])
gi3 = Input(shape = [1])
df = D(G([gi, gi2, gi3]))
adgan = Model(inputs = [gi, gi2, gi3], outputs = df)

lossとoptimizer

lossはGANには珍しくMSEを使った。

optimizerはAdamで学習率(lr)は下のように設定。

・Disgan => lr=0.0002

・Adage => lr=0.0001

両方0.0001よりdisganを少し大きくした方が精度がいい結果になった。

def gradient_penalty_loss(y_true, y_pred, averaged_samples, weight):
    gradients = K.gradients(y_pred, averaged_samples)[0]
    gradients_sqr = K.square(gradients)
    gradient_penalty = K.sum(gradients_sqr,
                              axis=np.arange(1, len(gradients_sqr.shape)))
    return K.mean(gradient_penalty * weight)

partial_gp_loss = partial(gradient_penalty_loss, averaged_samples = ri, weight = 5)

disgan.compile(optimizer=Adam(lr=0.0002, beta_1 = 0, beta_2 = 0.99, decay = 0.00001), loss=['mse', 'mse', partial_gp_loss])

adgan.compile(optimizer = Adam(lr=0.0001, beta_1 = 0, beta_2 = 0.99, decay = 0.00001), loss = 'mse')

3.訓練

主にdisganとadganのtrain。
従来のGANと訓練方法は同じだけど、generatorが違う分、disganとadgan共に独特な値を入れてる。

# Noise 
def noise(n):
    return np.random.normal(0.0, 1.0, size = [n, latent_size])

def noiseImage(n):
    return np.random.uniform(0.0, 1.0, size = [n, im_size, im_size, 1])

# train disgan
real_images  = x_train[idx*batch_size:(idx+1)*batch_size]
train_data = [real_images, noise(batch_size), noiseImage(batch_size), np.ones((batch_size, 1))]
d_loss = disgan.train_on_batch(train_data, [np.ones((batch_size, 1)),
                                                      -np.ones((batch_size, 1)), np.ones((batch_size, 1))])

# train adgan
g_loss = adgan.train_on_batch([noise(batch_size), noiseImage(batch_size), np.ones((batch_size, 1))], 
                                        np.zeros((batch_size, 1), dtype=np.float32))