Flux.jlがバージョンアップし、Juliaで機械学習：Flux.jlで自由自在にオリジナルレイヤーを組んでみようの内容が古くなってきていますので、新しく記事を書くことにしました。書き方を、最新のFlux.jlのマニュアルに書かれているような形に書き換えています。

はじめに

Juliaで機械学習するならFlux.jlが有名です。
https://github.com/FluxML/Flux.jl

これまで、いろいろを記事にまとめてきました。
Juliaで機械学習：深層学習フレームワークFlux.jlを使ってみる　その１：基本編
 Juliaで機械学習：深層学習フレームワークFlux.jlを使ってみる　その２：線形回帰編
 Juliaで機械学習：深層学習フレームワークFlux.jlを使ってみる　その3：ニューラルネットとバッチ正規化編
 Juliaで機械学習：Flux.jlで自由自在にオリジナルレイヤーを組んでみよう

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

環境

Julia 1.9.3
Flux 0.14.4

対象とする問題

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

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

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

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()

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

インプットデータ

まずインプットデータを作ります。
訓練用のデータとテスト用のデータを作っておきます。
これは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

と作ってしまいます。

これらをまとめると、

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は配列の要素にして入れてしまいましょう。

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)

テストデータも同様に

input_test,output_test = make_inputoutput(xtest,ytest,ztest)

にします。

次に、データの取り出しについてです。これまでの記事では自前でデータを取り出す関数を作っていましたが、今回はFlux.jlの機能を使ってみます。

train_loader = Flux.DataLoader((input_train,output_train), batchsize=5, shuffle=true);
for (x, y) in train_loader
    println(x)
    println(y)
end

このようにDataLoaderを使うと、データをいい感じにシャッフルして取り出してくれます。

モデルの設定

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

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

for layer in model
    display(layer.weight)
    display(layer.bias)
end

とします。これで、modelの中の各レイヤーの重みとバイアスを出力することができます。

次に、学習に使う最適化方法を設定します。これは、

rule = Adam()
opt_state = Flux.setup(rule, model)

で設定できます。

次に、loss関数です。最新版のFlux.jlでは、

loss(y_hat, y) = sum((y_hat .- y).^2)

のように定義します。このy_hatには、実際にはmodel(x)などが入ることになります。
loss関数の微分は

grads = Flux.gradient(m -> loss(m(input_train[:,1]), output_train[:,1]), model)[1]
display(grads)

で計算できます。そして、この微分を使ってモデルを1回更新するには、

Flux.update!(opt_state, model, grads)

とします。

訓練

次に訓練をする関数を作ります。

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

データローダを使うことでシンプルになりました。作ったモデルによる計算結果がこちらです。

全体のコード

ここまでの全体コードを以下に示します。

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に入っているとは限りませんので、独自のレイヤーを作ることもあるでしょう。

独自のレイヤーとして、

f(x) = \exp( W x + b)

という指数関数のレイヤーを作ることにします。学習したいパラメータはWとbです。そして、インプットの長さをk、アウトプットの長さを1とします。

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

struct Explayer{T,V}
    W::T
    b::V
end

そして、、訓練可能なパラメータを指定し、初期化用の関数を用意します。

Flux.trainable(a::Explayer) = (W=a.W,b=a.b)
Explayer(in::Integer) = Explayer(randn(1, in), randn(1))

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

function (m::Explayer)(x)
    exp.(m.W*x.+m.b)
end

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

Flux.@functor Explayer

登録します。

これで、これまでのDenseのようにChainの中で使える関数ができました。例えば、

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

とすれば、これまでと同様にモデルを訓練することができます。

より複雑なモデルを作成する。

Flux.jlでは、ResNetと呼ばれるモデルも簡単に作ることができます。
これは、何層か前の層での出力を別の層で受け取るようなものです。詳しく知りたい方はResNetなどの検索語で調べてみてください。
ResNetを作るにはParallelを使います。つまり、

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)の出力を返しており、それらを足します。

他の例は、前の記事を参考にしてみてください。