概要

ゼロから作るDeep Learningを元に機械学習ライブラリを使わずにGAN(Generative Adversarial Networks)の実装をして、手書き数字画像(MNIST)の生成をしました。
単純に真似するだけでは面白くないので、PyTorchっぽく作るようにしました。

コード

https://github.com/ObaTakeshi/ML_scratch
GANの他に、MNISTの分類の例もあげています。

注意

コードが拙いです。特に、型を指定していないことや、パラメータを引数にせず関数の中に直張りしている箇所があります。
正常に動かない箇所があります。(Batch Norm)
CNNなど実装していないものがたくさんあります。

発展・改良していただける人歓迎しています。

参考にしたもの

言語

Julia - v1.0.0

実装した機能

Layer関連
- Affine
- Sigmoid
- ReLU
- Leaky ReLU
- tanh
- Drop out
- Batch Normalization (進行中)
Optimizer関連
- SGD
- Adam

実装上の苦労

誤差逆伝播と計算ノードの理解が足りていなかったため、特に参考に載っていないtanh()などの逆伝播の実装に時間がかかりました。
実装する上で理解が進んだので、よかったです。
そして、実装中のBatch Normalizationの挙動が怪くて、現在苦労しています。

GANをうごかしてみる

ここでは、コードの代表的な部分を説明します。

読み込み

using Statistics
using Random
using Dates

include("./Functional.jl")
include("./Layer.jl")
include("./Optimizer.jl")
include("./MNIST.jl")

ここで、Functional.jlは、reluなど各種関数
Layer.jlは、Affine Layerなど各種Layer
Optimizer.jlは、Adam, SGDなど最適化手法
MNIST.jlは、MNISTデータ関連の操作を行うモジュールです。

データローダ

function dataloader(x, y, ;batch_size=1, shuffle=false)
    function producer(c::Channel, x, y, batch_size, shuffle)
        data_size = size(x, 2)
        if shuffle
            randidx = randperm(data_size)
            x = x[:, randidx]
            y = y[:, randidx]
        end
        i = 1
        while i < data_size-batch_size
            put!(c, (x[:, i:i+batch_size-1], y[:, i:i+batch_size-1]))
            i += batch_size
        end
        put!(c, (x[:, i:end], y[:, i:end]))
    end

    ch = Channel((ch_arg) -> producer(ch_arg, x, y, batch_size,  shuffle))
    return ch
end

dataloader()は、データとラベル(onehot)を受け取り、バッチサイズに切り分けて返すイテレータを作成しています。
ここで注意すべき点は、最後の次元がバッチサイズになる点です。

モデル定義

Generatorだけ説明します。

mutable struct Generator{T}
    a1lyr::Layer.AffineLayer{T}
    leakyrelu1lyr::Layer.ReluLayer
    a2lyr::Layer.AffineLayer{T}
    leakyrelu2lyr::Layer.ReluLayer
    a3lyr::Layer.AffineLayer{T}
    tanhlyr::Layer.TanhLayer
    params
end

function (::Type{Generator{T}})(input_size::Int, hidden_size::Int, hidden_size2::Int, output_size::Int; weight_init_std::Float64=0.1) where T
    W1 = weight_init_std .* randn(T, hidden_size, input_size)
    b1 = zeros(T, hidden_size)
    W2 = weight_init_std .* randn(T, hidden_size2, hidden_size)
    b2 = zeros(T, hidden_size2)
    W3 = weight_init_std .* randn(T, output_size, hidden_size2)
    b3 = zeros(T, output_size)

    a1lyr = Layer.AffineLayer(W1, b1)
    leakyrelu1lyr = Layer.ReluLayer()
    a2lyr = Layer.AffineLayer(W2, b2)
    leakyrelu2lyr = Layer.ReluLayer()
    a3lyr = Layer.AffineLayer(W3, b3)
    tanhlyr = Layer.TanhLayer()
    params = [a1lyr.W, a1lyr.b, a2lyr.W, a2lyr.b, a3lyr.W, a3lyr.b] #, bn1.gamma, bn1.beta, bn2.gamma, bn2.beta]
    Generator(a1lyr, leakyrelu1lyr, a2lyr, leakyrelu2lyr, a3lyr, tanhlyr, params)
end

function setparams(net::Generator, params)
    net.a1lyr.W = params[1]
    net.a1lyr.b = params[2]
    net.a2lyr.W = params[3]
    net.a2lyr.b = params[4]
    net.a3lyr.W = params[5]
    net.a3lyr.b = params[6]
end

function forward(net::Generator, x)
    x = Layer.forward(net.a1lyr, x)
    x = Layer.forward(net.leakyrelu1lyr, x)

    x = Layer.forward(net.a2lyr, x)
    x = Layer.forward(net.leakyrelu2lyr, x)

    x = Layer.forward(net.a3lyr, x)
    output = Layer.forward(net.tanhlyr, x)
    return output
end

function backward(net::Generator, y)
    y = Layer.backward(net.tanhlyr, y)
    y = Layer.backward(net.a3lyr, y)

    y = Layer.backward(net.leakyrelu2lyr, y)
    y = Layer.backward(net.a2lyr, y)

    y = Layer.backward(net.leakyrelu1lyr, y)
    y = Layer.backward(net.a1lyr, y)
    return [net.a1lyr.dW, net.a1lyr.db, net.a2lyr.dW, net.a2lyr.db, net.a3lyr.dW, net.a3lyr.db]
end

Julia のススメ〜 Deep Learning のための Julia 〜を参考に、型システムを用いてモデルやレイヤーを定義しています。
setparams()など冗長に見える関数がありますが、これは後述の最適化手法が返す重みを代入するために存在します。
forward(), backward()はその名の通り、順伝播、逆伝播です。

ハイパーパラメータなどの定義

const epochs = 100
const batch_size = 100
const learning_rate = Float64(1e-4)
const train_size = size(x_train, 2) # => 60000
const iter_per_epoch = Int32(max(train_size / batch_size, 1))
const noise_size = 100
const image_size = 28

const fixed_noise = randn(noise_size, 1)

generator = Generator{Float64}(100, 256, 512, 784)
discriminator = Discriminator{Float64}(784, 256, 128, 1)

gen_optimizer = Optimizer.Adam(generator, learning_rate)
dis_optimizer = Optimizer.SGD(discriminator, learning_rate)

ここでは、epoch数などのハイパーパラメータの定義の他に、Generatorなどのネットワーク、最適化手法の定義を行なっています。
Generator()やDiscriminator()の引数は、入力層・隠れ層1・隠れ層2・出力層のユニット数です。
今回、GeneratorにはAdam、DiscriminatorにはSGDを用います。(現在の実装では、AdamとSGDの組み合わせで一番良い生成ができるようです。)

学習

for epoch = 1:epochs
    iter = 0
    for (x_batch, _) in dataloader(x_train, y_train, batch_size=batch_size, shuffle=true)
        iter += 1
        start_time = now()
        # train Real
        label = ones(1, batch_size)
        d_x = output = forward(discriminator, x_batch)
        dr_loss = criterion(discriminator, output, label)
        grads = backward(discriminator, dr_loss)
        Optimizer.step(dis_optimizer, grads)

        # train fake
        noise = randn(noise_size, batch_size)
        fake = forward(generator, noise)

        label = zeros(1, batch_size)
        d_g = output = forward(discriminator, fake)
        df_loss = criterion(discriminator, output, label)
        grads = backward(discriminator, df_loss)
        Optimizer.step(dis_optimizer, grads)

        # train Generator
        label = ones(1, batch_size)
        output = forward(discriminator, fake)
        g_loss = criterion(discriminator, output, label)
        y = _backward(discriminator, g_loss)
        grads = backward(generator, y)
        Optimizer.step(gen_optimizer, grads)

        d_loss = dr_loss + df_loss
        if iter % 100 == 0
            fake = forward(generator, fixed_noise)
            save_checkpoints(fake, epoch, iter)
            print("D(x): $(mean(d_x)) ")
            println("D(G(z)): $(mean(d_g))")
            println("epoch: [$(epoch)/$(epochs)] [$(iter)/$(iter_per_epoch)] g_loss: $(g_loss) d_loss: $(d_loss)")
        end
    end
    fake = forward(generator, fixed_noise)
    save_checkpoints(fake, epoch, iter)
end

学習のコードをなるべくPyTorchライクになるよう実装しました。
学習の流れとしては、Discriminatorの学習(real → fake)、Generatorの学習の順です。
中では、forward()で順伝播、criterion()でロスの計算、backward()で逆伝播、最後にOptimizer.step()で重みの更新という流れです。

Generator学習時に、_backward()という行があります。これは、Generatorの逆伝播はDiscriminatorから一連の流れとして続くため、Discriminator → Generatorという風に逆伝播を通すために、Discriminatorの定義に追加しています。

save_checkpoints()は、現在のところ生成画像のする関数です。

GANによる生成画像

1epoch目

5epoch目

10epoch目

25epoch目

今後の展望

型指定やinline化などを行う
ConvolutionなどCNN関連、RNN関連を実装する
Batch Normalizationや正則化などの実装を行う

興味を持った方など、ぜひプルリクエストなどしていただけるとありがたいです。

Juliaで機械学習をゼロから作りGANを試してみる

概要

コード

注意