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

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

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

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

## 2023年9月追記
Fluxのバージョンアップに合わせて新しい記事[Juliaで機械学習：Flux.jlで自由自在にオリジナルレイヤーを組んでみよう 2023年版](https://qiita.com/cometscome_phys/items/63418f73e84afbdc2d4c)も書きましたので、そちらも参考にしてください。


# 環境
Julia 1.4.2
Flux 0.11.0

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

```math
f(x,y) = xy + \cos(3x)+xe^{y/5} + \tanh(y) \cos(3y)
```
とします。2次元平面上での値を出す関数ですね。
これは

```julia
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")
```

![original.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/246113/b27de0cf-adbd-0477-7392-236a7701a973.png)

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

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

```julia
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のデータを作りました。次に、教師データを

```julia
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
```
と作ってしまいます。

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

```julia
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
```

テストデータも同様に

```julia
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
```
にします。

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

```julia
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
```
という関数を用意しました。ここまでは良くある形ですね。


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

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

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

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

最適化方法をADAM

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


```julia
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
```
となりますので、これを実行すれば学習ができます。

```julia
nt = 3000
train_batch!(inputdata_train,inputdata_test,model,loss,opt,nt) #学習
```
これでニューラルネットの学習ができました。簡単ですね！

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

![dense.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/246113/c078e56e-d094-6ddb-6395-035033798970.png)

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

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

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

```julia
struct Dense2
    W
    b
    activation
end
```
ここではDense2というレイヤーを定義しました。
このままではただのstructなので、

```julia
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と同じですね。

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

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

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

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

例えば、

```julia
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を入れてみます。
以下のようにします。

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

```julia
inner = Innerchain(Chain(Dense2(10,10,relu),Dense2(10,10,relu)))
model = Chain(Dense2(2,10,relu),inner,Dense2(10,1))
```
このように、modelのChainの中にChainを入れることができます。
これができると何が嬉しいか、わかりにくいかもしれませんので、次のものを見てみましょう。

## スキップコネクション
以下のように新しいレイヤーを定義してみます。

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

```julia
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を足す、ということができています。

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

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

```julia
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")
```

![denseskip.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/246113/d69088b5-01ae-f845-b4e3-b3565507197a.png)
という結果が得られます。

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

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

トレーニング用：

```julia
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)
```

テスト用：

```julia
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を独立にインプットとして入力して、最後に足すようなニューラルネットワークを考えてみましょう。

```math
f(x,y) = f_x(x) + f_y(y)
```
これは今回の問題では意味はありませんが、こんなネットワークを作りたいこともあるでしょう。
新しいレイヤーとしてDensexyを以下のように作ります。

```julia
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を作ると、例えば、

```julia
model = Chain(Densexy(1,10,σ),Densexy(10,10,σ),Densexy(10,10,σ),Densexy(10,1))
```
となります。
これの出力結果は

```julia
model(inputdata_train[1][1])
```

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

つまり、Chainで、

```julia
model = Chain(Densexy(1,10,σ),Densexy(10,10,σ),Densexy(10,10,σ),Densexy(10,1),x-> sum(x))
```
とすれば、最後の出力に対して値を変更できます。ここでは和を取っていますね。
これは

```julia
model = Chain(Densexy(1,10,σ),Densexy(10,10,σ),Densexy(10,10,σ),Densexy(10,1),sum)
```
とも書けます。結果は同じです。
この場合は、定義されているsumという関数を使っています。

あとは同じように

```julia
nt = 3000
train_batch!(inputdata_train,inputdata_test,model,loss,opt,nt) #学習
```
でトレーニングできます。なお、結果は載せていませんが、全然合わないです。

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

```math
x \cos(x)+\cosh(x)+\sin(2x)+ y^2 \tanh(y) + \exp(y/5)\cos(\sqrt{|y|})) 
```
としました。
これをプロットすると、

```julia
    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")
```

![originalde.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/246113/a7603f01-27a7-a49b-f29d-8a4f02150327.png)
となります。
この関数はxに関する関数とyに関する関数の和となっています。ですので、二つのニューラルネットワークをつなげるような形にしてもフィットできるはずです。それをやってみます。

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

```julia
    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
```

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

```julia
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で足し合わせるようになっています。
そして、

```julia
    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は同じです。

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

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

得られる結果は、

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

![dense2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/246113/504954a2-526a-a214-e902-f5707b2e3af3.png)
となり、バッチリですね。