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

Flux.jlがバージョンアップし、[Juliaで機械学習：Flux.jlで自由自在にオリジナルレイヤーを組んでみよう ](https://qiita.com/cometscome_phys/items/7452bda576a4bf26c1d1)の内容が古くなってきていますので、新しく記事を書くことにしました。書き方を、最新の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.jlで自由自在にオリジナルレイヤーを組んでみよう ](https://qiita.com/cometscome_phys/items/7452bda576a4bf26c1d1)

今回は2023年9月現在でFlux.jlを使う場合についてまとめておこうと思います。

# 環境
Julia 1.9.3
Flux 0.14.4

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

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

```julia
using Plots
function main()
    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")
end
main()
```

![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
function make_data(f)
    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
    return xtrain,ytrain,ztrain,xtest,ytest,ztest
end
xtrain,ytrain,ztrain,xtest,ytest,ztest = make_data(f)
```

となります。

そして、訓練データを作ります。インプットの配列とアウトプットの配列を作ることにして、x,yは配列の要素にして入れてしまいましょう。

```julia
function make_inputoutput(x,y,z)
    count = 0
    numx = length(x)
    numy = length(y)
    input = zeros(Float64,2,numx*numy)
    output = zeros(Float64,2,numx*numy)
    count = 0
    for i=1:numx
        for j=1:numy
            count += 1
            input[1,count] = x[i]
            input[2,count] = y[j]
            output[1,count] = z[count]
        end
    end
    return input,output
end

input_train,output_train = make_inputoutput(xtrain,ytrain,ztrain)

```

テストデータも同様に

```julia
input_test,output_test = make_inputoutput(xtest,ytest,ztest)
```
にします。

次に、データの取り出しについてです。これまでの記事では自前でデータを取り出す関数を作っていましたが、今回はFlux.jlの機能を使ってみます。
```julia
train_loader = Flux.DataLoader((input_train,output_train), batchsize=5, shuffle=true);
for (x, y) in train_loader
    println(x)
    println(y)
end
```
このように```DataLoader```を使うと、データをいい感じにシャッフルして取り出してくれます。

## モデルの設定

```julia
model = Chain(Dense(2,10,relu),Dense(10,10,relu),Dense(10,10,relu),Dense(10,1))
```
これで、密に結合したニューラルネットワークができました。
インプットはxとyで次元が2、次に隠れ層を3層用意して、それぞれのユニットが10あります。
その隠れ層の活性化関数はReLuです。
このモデルの中身を見たければ、
```julia
for layer in model
    display(layer.weight)
    display(layer.bias)
end
```
とします。これで、modelの中の各レイヤーの重みとバイアスを出力することができます。

次に、学習に使う最適化方法を設定します。これは、
```julia
rule = Adam()
opt_state = Flux.setup(rule, model)
```
で設定できます。

次に、loss関数です。最新版のFlux.jlでは、
```julia
loss(y_hat, y) = sum((y_hat .- y).^2)
```
のように定義します。この```y_hat```には、実際には```model(x)```などが入ることになります。
loss関数の微分は
```julia
grads = Flux.gradient(m -> loss(m(input_train[:,1]), output_train[:,1]), model)[1]
display(grads)
```
で計算できます。そして、この微分を使ってモデルを1回更新するには、
```julia
Flux.update!(opt_state, model, grads)
```
とします。

## 訓練
次に訓練をする関数を作ります。
```julia
function train_batch!(x_train,y_train,model,loss,opt_state,x_test,y_test,nepoch,batchsize)
    numtestdata = size(y_test)[2]
    train_loader = Flux.DataLoader((x_train,y_train), batchsize=batchsize, shuffle=true);
    for it=1:nepoch
        for (x, y) in train_loader
            grads = Flux.gradient(m -> loss(m(x), y), model)[1]
            Flux.update!(opt_state, model, grads)
        end

        if it % 10 == 0
            lossvalue = loss(model(x_test),y_test)/numtestdata 
            println("$it-th testloss: $lossvalue")
        end
    end
end
```
データローダを使うことでシンプルになりました。作ったモデルによる計算結果がこちらです。

![dense.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/246113/dcd97921-37f3-5fe1-9a5c-e99b19c73e6b.png)

## 全体のコード
ここまでの全体コードを以下に示します。
```julia
using Plots
using Flux

function make_data(f)
    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
    return xtrain,ytrain,ztrain,xtest,ytest,ztest
end

function make_inputoutput(x,y,z)
    count = 0
    numx = length(x)
    numy = length(y)
    input = zeros(Float64,2,numx*numy)
    output = zeros(Float64,2,numx*numy)
    count = 0
    for i=1:numx
        for j=1:numy
            count += 1
            input[1,count] = x[i]
            input[2,count] = y[j]
            output[1,count] = z[count]
        end
    end
    return input,output
end


function train_batch!(x_train,y_train,model,loss,opt_state,x_test,y_test,nepoch,batchsize)
    numtestdata = size(y_test)[2]
    train_loader = Flux.DataLoader((x_train,y_train), batchsize=batchsize, shuffle=true);
    for it=1:nepoch
        for (x, y) in train_loader
            grads = Flux.gradient(m -> loss(m(x), y), model)[1]
            Flux.update!(opt_state, model, grads)
        end

        if it % 10 == 0
            lossvalue = loss(model(x_test),y_test)/numtestdata 
            println("$it-th testloss: $lossvalue")
        end
    end
end


function main()
    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")

    xtrain,ytrain,ztrain,xtest,ytest,ztest = make_data(f)

    input_train,output_train = make_inputoutput(xtrain,ytrain,ztrain)
    input_test,output_test = make_inputoutput(xtest,ytest,ztest)

    #train_loader = Flux.DataLoader((input_train,output_train), batchsize=5, shuffle=true);
    model = Chain(Dense(2,10,relu),Dense(10,10,relu),Dense(10,10,relu),Dense(10,1))
    rule = Adam()
    opt_state = Flux.setup(rule, model)
    loss(y_hat, y) = sum((y_hat .- y).^2)
    nepoch = 1000
    batchsize = 128
    train_batch!(input_train,output_train ,model,loss,opt_state,input_test,output_test,nepoch,batchsize)

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

end
main()

```

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

独自のレイヤーとして、
```math
f(x) = \exp( W x + b)
```
という指数関数のレイヤーを作ることにします。学習したいパラメータはWとbです。そして、インプットの長さをk、アウトプットの長さを1とします。

まず、独自レイヤーに名前をつけ、その中に入っているパラメータを定義します。
```julia
struct Explayer{T,V}
    W::T
    b::V
end
```
そして、、訓練可能なパラメータを指定し、初期化用の関数を用意します。
```julia
Flux.trainable(a::Explayer) = (W=a.W,b=a.b)
Explayer(in::Integer) = Explayer(randn(1, in), randn(1))
```
次に、レイヤーの挙動を記述します。
```julia
function (m::Explayer)(x)
    exp.(m.W*x.+m.b)
end
```

最後に、ここで作ったstructがFluxで使えるということを
```julia
Flux.@functor Explayer
```
登録します。

これで、これまでのDenseのようにChainの中で使える関数ができました。例えば、
```julia
model = Chain(Dense(2,10,relu),Dense(10,10,relu),Dense(10,10,relu),Explayer(10))
```
とすれば、これまでと同様にモデルを訓練することができます。

## より複雑なモデルを作成する。
Flux.jlでは、ResNetと呼ばれるモデルも簡単に作ることができます。
これは、何層か前の層での出力を別の層で受け取るようなものです。詳しく知りたい方はResNetなどの検索語で調べてみてください。
ResNetを作るには```Parallel```を使います。つまり、
```julia
resnet = Parallel(+,x->x,Dense(10,10,relu),Dense(10,10,relu))
model = Chain(Dense(2,10,relu),resnet,Explayer(10))
```
のようにします。ここで、```Parallel(f,a,b)```のような形をしていますが、これはaの出力とbの出力を用いて```f(a,b)```とする、という意味です。上の例では、aの出力とbの出力を足しています。aはxをそのまま、bは```Dense(10,10,relu),Dense(10,10,relu)```の出力を返しており、それらを足します。

他の例は、[前の記事](https://qiita.com/cometscome_phys/items/7452bda576a4bf26c1d1)を参考にしてみてください。