Juliaで機械学習：Flux.jlで自由自在にオリジナルレイヤーを組んでみよう

Juliaで機械学習をするならFlux.jlが有名ですね。
https://github.com/FluxML/Flux.jl
しかしドキュメントがイマイチわかりにくいような気がします。ということで、以前から色々まとめてきました。

Juliaで機械学習：深層学習フレームワークFlux.jlを使ってみる　その１：基本編
Juliaで機械学習：深層学習フレームワークFlux.jlを使ってみる　その２：線形回帰編
Juliaで機械学習：深層学習フレームワークFlux.jlを使ってみる　その3：ニューラルネットとバッチ正規化編
しかし、JuliaもFluxもバージョンアップしていて、上に書いたコードがそのまま動かない例が出てきました。
ということで、この辺りで綺麗にまとめてみようと思います。

特に、自分で好きなようにニューラルネットワークを組みたい！という時にどうやるかについて述べます。例えばResNetのskip connectionとか、あるいは普通とは異なる重みを持つレイヤーをつけるとか。
今回書いたやり方を使えば、原理的にはどんなニューラルネットも作ることができそうです。

2023年9月追記

Fluxのバージョンアップに合わせて新しい記事Juliaで機械学習：Flux.jlで自由自在にオリジナルレイヤーを組んでみよう 2023年版も書きましたので、そちらも参考にしてください。

環境

Julia 1.4.2
Flux 0.11.0

対象とする問題

何か具体的な問題を通じて考えた方がわかりやすいと思いますので、今回もこれまでと同様に関数をフィッティングすることにします。MNISTなどがやりたい場合は適宜インプットとアウトプットを読み替えてください。
考える式は、インプットをx,yとして

f(x,y) = xy + \cos(3x)+xe^{y/5} + \tanh(y) \cos(3y)

とします。2次元平面上での値を出す関数ですね。
これは

using Plots
num = 30
x = range(-2,2,length=num)
y = range(-2,2,length=num)
f(x,y) = x*y + cos(3*x)+exp(y/5)*x + tanh(y)*cos(3*y)
z =[f(i,j) for i in x, j in y]'
p = plot(x,y,z, st=:wireframe)
savefig("original.png")

こんな感じになります。
この関数をフィッティングしてみましょう。

インプットデータ

まずインプットデータを作ります。
訓練用のデータとテスト用のデータを作っておきます。
これはx,yの組を作れば良いわけですね。なので、

num = 47
numt = 19
numtrain = num*num
numtest = numt*numt
xtrain = range(-2,2,length=num)
ytrain = range(-2,2,length=num)
xtest = range(-2,2,length=numt)
ytest = range(-2,2,length=numt)

とxとyのデータを作りました。次に、教師データを

count = 0
ztrain = Float32[]
for i = 1:num
    for j=1:num
        count += 1
        push!(ztrain, f(xtrain[i],ytrain[j]))
    end
end

count = 0
ztest = Float32[]
for i = 1:numt
    for j=1:numt
        count += 1
        push!(ztest, f(xtest[i],ytest[j]))
    end
end

と作ってしまいます。

そして、訓練データを(インプット,アウトプット)というタプルの組にして作ります。

inputdata_train = []
count = 0
for i=1:num
    for j=1:num
        count += 1
        push!(inputdata_train,([xtrain[i],ytrain[j]],ztrain[count]))
    end
end

テストデータも同様に

inputdata_test = []
count = 0
for i=1:numt
    for j=1:numt
        count += 1
        push!(inputdata_test,([xtest[i],ytest[j]],ztest[count]))
    end
end

にします。

また、ランダムバッチを作成するために、

function make_random_batch(data_input,batchsize)
    numofdata = length(data_input)
    A = rand(1:numofdata,batchsize) #インデックスをシャッフル
    data = []
    for i=1:batchsize
        push!(data,data_input[A[i]]) #ランダムバッチを作成。 
    end
    return data
end

という関数を用意しました。ここまでは良くある形ですね。

モデルの作成

では、モデルを作っていきましょう。まず、これをみてください。

model = Chain(Dense(2,10,relu),Dense(10,10,relu),Dense(10,10,relu),Dense(10,1))

これで、密に結合したニューラルネットワークができました。
インプットはxとyで次元が2、次に隠れ層を3層用意して、それぞれのユニットが10あります。
その隠れ層の活性化関数はReLuです。
はい、これで終わりです。簡単ですね。

これを用いてloss関数を平均自乗誤差

loss(x,y) = Flux.mse(model(x), y)

最適化方法をADAM

opt = ADAM() 

とすれば、あとはトレーニングするだけです。
トレーニングは
Flux.train!(loss, params(model),data, opt)を呼べば良いだけです。
dataをランダムバッチで作るとすると、

function train_batch!(data_train,data_test,model,loss,opt,nt)
    batchsize = 128
    for it=1:nt
        data = make_random_batch(data_train,batchsize)
        Flux.train!(loss, params(model),data, opt)
        if it% 100 == 0
            lossvalue = 0.0
            #testmode!(model, true)
            for i=1:length(data_test)                
                lossvalue += loss(data_test[i][1],data_test[i][2])
            end
            #testmode!(model, false)
            println("$(it)-th loss = ",lossvalue/length(data_test))
        end
    end
end

となりますので、これを実行すれば学習ができます。

nt = 3000
train_batch!(inputdata_train,inputdata_test,model,loss,opt,nt) #学習

これでニューラルネットの学習ができました。簡単ですね！

znn =[model([i,j])[1] for i in x, j in y]'
p = plot(x,y,[znn], st=:wireframe)
savefig("dense.png")

ちょっとオリジナルよりカクカクしている感じはありますが、だいたい合ってます。

自分でニューラルネットワークを作りたい

さて、上のように密な結合のニューラルネットワークを作るときは簡単にできることがわかりました、また、Fluxに入っている別のレイヤーを使えば、同様に作ることができそうです。しかし、自分がやりたいことが全てFluxに入っているとは限りませんので、独自のレイヤーを作ることもあるでしょう。

独自のレイヤーを作るには、以下のようにします。
まず、独自レイヤーに名前をつけ、その中に入っているパラメータを定義します。

struct Dense2
    W
    b
    activation
end

ここではDense2というレイヤーを定義しました。
このままではただのstructなので、

Flux.trainable(a::Dense2) = (a.W,a.b)
Dense2(in::Integer, out::Integer) = Dense2(randn(out, in), randn(out),x -> x)
Dense2(in::Integer, out::Integer,activation) = Dense2(randn(out, in), randn(out),activation)

とします。1行目はパラメータとしてトレーニングできるものを指定しています。activationは活性化関数を入れたいので今はtrainableはWとbだけです。
次の2行目は、Dense2のコンストラクタで、三番目の引数がないときは恒等操作x->xをactivationに入れています。
3行目はactivationがあるときですね。
inとかoutとかで層のユニットの数を指定しています。Denseと同じですね。

次に、レイヤーの挙動を記述します。

function (m::Dense2)(x)
    m.activation.(m.W * x .+ m.b)
end

これは、よく見る形ですね。W*x+bをしてから、活性化関数をかけています。
つまり、インプットxに対してアウトプットがどうなるか、ということを書いています。
なお、xはベクトルでも行列でもTupleでも良いです。ともかく、ここに書いた演算が実行可能であり、アウトプットがあれば良いです。何度もDense2をかけるのであれば、インプットのxと同じようなアウトプットの方が連結できて良いですね。

最後に、ここで作ったstructがFluxで使えるということを

Flux.@functor Dense2

で指定しておきます。
これでFluxのtrain!で使えるレイヤーができました。

例えば、

model = Chain(Dense2(2,10,relu),Dense2(10,10,relu),Dense2(10,10,relu),Dense2(10,1))

このようにDense2というものをChainの中で使うことができます。

スキップコネクション

次に、スキップコネクションを実装してみましょう。
これは、何層か前の層での出力を別の層で受け取るようなものです。詳しく知りたい方はResNetなどの検索語で調べてみてください。
スキップコネクションを実装するのは、Chainを使った素朴なやり方で難しそうに見えます。
しかし、そんなことはありません。Chainの中にChainを入れてあげれば良いのです。

Chainの中のChain

まず、Chainの中にChainを入れてみます。
以下のようにします。

struct Innerchain
    inner
end
function (m::Innerchain)(x)
    m.inner(x)
end    
Flux.@functor  Innerchain

新しいレイヤーであるInnerchainというものを定義しました。このInnerchainの挙動を記述していますが、structのフィールドであるinnerを使ってインプットxからアウトプットxを作っています。
これを用いると、

inner = Innerchain(Chain(Dense2(10,10,relu),Dense2(10,10,relu)))
model = Chain(Dense2(2,10,relu),inner,Dense2(10,1))

このように、modelのChainの中にChainを入れることができます。
これができると何が嬉しいか、わかりにくいかもしれませんので、次のものを見てみましょう。

スキップコネクション

以下のように新しいレイヤーを定義してみます。

struct IdensitySkip
    inner
end
function (m::IdensitySkip)(x)
    m.inner(x) .+ x
end
Flux.@functor  IdensitySkip

IdentitySkipが新しいレイヤーです。structの定義は先ほどと変わりません。
変わったのは挙動の部分です。inner(x)をやった後にxを足したものをアウトプットにしています。
つまり、このレイヤーに入ってきたインプットxを、このレイヤーの中のレイヤーの出力結果に足しています。
innerは複数のレイヤーで構成されるわけですから、これは複数のレイヤーをスキップしてつなげたことになります。

inner = IdensitySkip(Chain(Dense2(10,10,relu),Dense2(10,10,relu)))
model = Chain(Dense2(2,10,relu),inner,Dense2(10,1))

とすれば、Dense2が2枚入っており、xにDense2を2回作用させ、その後にxを足す、ということができています。

もちろんちゃんと学習ができて、

nt = 3000
train_batch!(inputdata_train,inputdata_test,model,loss,opt,nt) #学習

znn =[model([i,j])[1] for i in x, j in y]'
ztrue =[f(i,j) for i in x, j in y]'
plot(x,y,znn,st=:wireframe)
p = plot(x,y,[znn], st=:wireframe)
savefig("denseskip.png")

という結果が得られます。

以上、好きなレイヤーを作ったりスキップコネクションを作ったりする方法でした。

追記

最後に、ちょっと特殊なニューラルネットワークを作ってみましょう。特に、ラムダ関数を使ってみます。
インプットデータは、これまではxとyの2次元の配列でしたが、今度はTupleにしてみます。

トレーニング用：

inputdata_train = []
count = 0
for i=1:num
    for j=1:num
        count += 1
        push!(inputdata_train,((xtrain[i],ytrain[j]),ztrain[count]))
    end
end

2209-element Array{Any,1}:
 ((-2.0, -2.0), 2.693899556678832)
 ((-2.0, -1.9130434782608696), 2.6029812341505876)
 ((-2.0, -1.826086956521739), 2.5661321750241433)
 ((-2.0, -1.7391304347826086), 2.5711460568736983)
 ((-2.0, -1.6521739130434783), 2.6027049817080266)
 ((-2.0, -1.565217391304348), 2.6435034080981814)
 ((-2.0, -1.4782608695652173), 2.67556153879426)
 ((-2.0, -1.391304347826087), 2.681642326233047)
 ((-2.0, -1.3043478260869565), 2.6466718363907957)
 ((-2.0, -1.2173913043478262), 2.5590522051958344)
 ((-2.0, -1.1304347826086956), 2.411750966304716)
 ((-2.0, -1.0434782608695652), 2.203051747345016)
 ((-2.0, -0.9565217391304348), 1.9368614874796148)
 ⋮
 ((2.0, 1.0434782608695652), 4.732057357525006)
 ((2.0, 1.1304347826086956), 4.942400614845289)
 ((2.0, 1.2173913043478262), 5.2144177221369805)
 ((2.0, 1.3043478260869565), 5.546413930940305)
 ((2.0, 1.391304347826087), 5.931378488805956)
 ((2.0, 1.4782608695652173), 6.3577353421055935)
 ((2.0, 1.565217391304348), 6.810432470404888)
 ((2.0, 1.6521739130434783), 7.272254849504222)
 ((2.0, 1.7391304347826086), 7.7252448219527725)
 ((2.0, 1.826086956521739), 8.15211910925267)
 ((2.0, 1.9130434782608696), 8.537582206209525)
 ((2.0, 2.0), 8.869450319833163)

テスト用：

inputdata_test = []
count = 0
for i=1:numt
    for j=1:numt
        count += 1
        push!(inputdata_test,((xtest[i],ytest[j]),ztest[count]))
    end
end

361-element Array{Any,1}:
 ((-2.0, -2.0), 2.693899556678832)
 ((-2.0, -1.7777777777777777), 2.5646702274917677)
 ((-2.0, -1.5555555555555556), 2.647823977942668)
 ((-2.0, -1.3333333333333333), 2.663690531543826)
 ((-2.0, -1.1111111111111112), 2.370630068456955)
 ((-2.0, -0.8888888888888888), 1.6958641005960973)
 ((-2.0, -0.6666666666666666), 0.7856802609979634)
 ((-2.0, -0.4444444444444444), -0.07900501416745546)
 ((-2.0, -0.2222222222222222), -0.6802652741398586)
 ((-2.0, 0.0), -1.0398297133496341)
 ((-2.0, 0.2222222222222222), -1.403345420193552)
 ((-2.0, 0.4444444444444444), -2.0164672893263416)
 ((-2.0, 0.6666666666666666), -2.9009479493745713)
 ⋮
 ((2.0, -0.4444444444444444), 1.803006122974891)
 ((2.0, -0.2222222222222222), 2.2569607933833695)
 ((2.0, 0.0), 2.960170286650366)
 ((2.0, 0.2222222222222222), 3.6673310475515053)
 ((2.0, 0.4444444444444444), 4.13314732712037)
 ((2.0, 0.6666666666666666), 4.3362419644749854)
 ((2.0, 0.8888888888888888), 4.494894833531509)
 ((2.0, 1.1111111111111112), 4.89037788145797)
 ((2.0, 1.3333333333333333), 5.669337042490845)
 ((2.0, 1.5555555555555556), 6.759355627817648)
 ((2.0, 1.7777777777777777), 7.919159829730134)
 ((2.0, 2.0), 8.869450319833163)

ここではタプルとしてxとyが入っていますが、このxとyがそれぞれ配列だったりするインプットもありえますね。
そして、xとyを独立にインプットとして入力して、最後に足すようなニューラルネットワークを考えてみましょう。

f(x,y) = f_x(x) + f_y(y)

これは今回の問題では意味はありませんが、こんなネットワークを作りたいこともあるでしょう。
新しいレイヤーとしてDensexyを以下のように作ります。

struct Densexy
    Wx
    bx
    Wy
    by
    activation
end

Flux.trainable(a::Densexy) = (a.Wx,a.bx,a.Wy,a.by)
Densexy(in::Integer, out::Integer) = Densexy(randn(out, in), randn(out),randn(out, in),randn(out),x -> x)
Densexy(in::Integer, out::Integer,activation) = Densexy(randn(out, in), randn(out),randn(out, in),randn(out),activation)

function (m::Densexy)(x)
    (m.activation.(m.Wx * x[1] .+ m.bx), m.activation.(m.Wy * x[2] .+ m.by) )
end
    
Flux.@functor Densexy

こうすると、xはタプルになっており、タプルが返ってきますね。また、パラメータは四つです。

この新しいレイヤーを使ってmodelを作ると、例えば、

model = Chain(Densexy(1,10,σ),Densexy(10,10,σ),Densexy(10,10,σ),Densexy(10,1))

となります。
これの出力結果は

model(inputdata_train[1][1])

([1.5053051197361837], [-2.0112338073840363])

となります。期待した通り、タプルが返ってきますね。
これでxとyを別々にインプットにしてニューラルネットワークに入れることができました。
最後に、この二つの和を取りたいですね。
Chainの中でどうやるか、ですが、無名関数が使えます。
無名関数というのはラムダ関数とも言いますね。いちいち関数を定義しないでその場で作るやつです。

つまり、Chainで、

model = Chain(Densexy(1,10,σ),Densexy(10,10,σ),Densexy(10,10,σ),Densexy(10,1),x-> sum(x))

とすれば、最後の出力に対して値を変更できます。ここでは和を取っていますね。
これは

model = Chain(Densexy(1,10,σ),Densexy(10,10,σ),Densexy(10,10,σ),Densexy(10,1),sum)

とも書けます。結果は同じです。
この場合は、定義されているsumという関数を使っています。

あとは同じように

nt = 3000
train_batch!(inputdata_train,inputdata_test,model,loss,opt,nt) #学習

でトレーニングできます。なお、結果は載せていませんが、全然合わないです。

追記2

別のニューラルネットワークも作ってみましょう。
今度は関数を

x \cos(x)+\cosh(x)+\sin(2x)+ y^2 \tanh(y) + \exp(y/5)\cos(\sqrt{|y|})) 

としました。
これをプロットすると、

    num = 30
    x = range(-2,2,length=num)
    y = range(-2,2,length=num)
    f(x,y) =x*cos(x)+cosh(x)+sin(2x)+ y^2*tanh(y) + exp(y/5)*cos(sqrt(abs(y))) 
    z =[f(i,j) for i in x, j in y]'
    p = plot(x,y,z, st=:wireframe)
    savefig("originalde.png")

となります。
この関数はxに関する関数とyに関する関数の和となっています。ですので、二つのニューラルネットワークをつなげるような形にしてもフィットできるはずです。それをやってみます。

先ほどと同じように訓練データとテストデータを作ります。

    num = 47
    numt = 19
    numtrain = num*num
    numtest = numt*numt
    xtrain = range(-2,2,length=num)
    ytrain = range(-2,2,length=num)
    xtest = range(-2,2,length=numt)
    ytest = range(-2,2,length=numt)

    count = 0
    ztrain = Float32[]
    for i = 1:num
        for j=1:num
            count += 1
            push!(ztrain, f(xtrain[i],ytrain[j]))
        end
    end

    count = 0
    ztest = Float32[]
    for i = 1:numt
        for j=1:numt
            count += 1
            push!(ztest, f(xtest[i],ytest[j]))
        end
    end

    inputdata_train = []
    count = 0
    for i=1:num
        for j=1:num
            count += 1
            push!(inputdata_train,([xtrain[i],ytrain[j]],ztrain[count]))
        end
    end

    inputdata_test = []
    count = 0
    for i=1:numt
        for j=1:numt
            count += 1
            push!(inputdata_test,([xtest[i],ytest[j]],ztest[count]))
        end
    end

複数のニューラルネットワークをつなげるには、

using Flux
struct BPchain
    inners
end

function (m::BPchain)(x)
    y = sum([m.inners[i]([x[i]]) for i=1:length(x)])
    return y
end    
Flux.@functor  BPchain

というFlux用のstructを作っておきます。BPchainでは、それぞれのニューラルネットワークでの評価の後にsumで足し合わせるようになっています。
そして、

    nmiddle = 10
    nparam = 1
    inner1 = Chain(Dense(nparam,nmiddle,tanh),Dense(nmiddle,nmiddle,tanh),Dense(nmiddle,1))
    inner2 = Chain(Dense(nparam,nmiddle,tanh),Dense(nmiddle,nmiddle,tanh),Dense(nmiddle,1))
    inners = (inner1,inner2)
    model = BPchain(inners)

のように、2種類のChainを作り、それをBPchainに入れ込みます。
loss関数やoptは同じです。

    loss(x,y) = Flux.mse(model(x), y)
    opt = ADAM() 

    nt = 3000
    train_batch!(inputdata_train,inputdata_test,model,loss,opt,nt) #学習

得られる結果は、

    znn =[model([i,j])[1] for i in x, j in y]'
    p = plot(x,y,znn, st=:wireframe)
    savefig("dense2.png")

となり、バッチリですね。