Julia 1.0 ＋ FluxでMNIST学習

# 注意

この記事はすでに古く，いくつか正常に動作しないことがわかっています．

# 本日は

Julia 0.6 から 0.7、 1.0 へ粛々と環境を移動しつつあります。

Ubuntu，Macであれば

- Images
- DataFrames
- IJulia
- PyCall
- PyPlot
- pyjulia(PythonからJuliaを呼ぶライブラリ)

は動いているように思えます。Windowsについては最近Windowsを動かすこと自体を仕事でもプライベートでもしていないので確認していません。

せっかくなのでDeepLearningしたい。Juliaでできるのはあるのでしょうか？

- [MXNet.jl](https://github.com/dmlc/MXNet.jl)
  - The current support julia version is 0.6. Julia 0.7/1.0 is not supported yet.
  - がーん・・・。[Hands-On-Computer-Vision-with-Julia](https://github.com/PacktPublishing/Hands-On-Computer-Vision-with-Julia)
 を動かそうと目論んでいたのですが、Julia1.0ではまだサポートされていないようです。

- [Merlin.jl](https://github.com/hshindo/Merlin.jl)
  - Merlin is tested against Julia 1.0 on Linux, OS X, and Windows (x64).
  - インストール、`using Merlin` は動きそう。ただ、動作確認のために mnist.jl がデータセットのローダー周りでエラーを起こします。この現象Julia0.6でもあったのでこの機会にIssueを書いておきました[Datasets not defined #20](https://github.com/hshindo/Merlin.jl/issues/20)。`This will be fixed soon` とのことです。

- TensorFlow.jl
 - インストール, using TensorFlow まではいけてます。 （ゲッソリ）
- [Flux.jl](https://github.com/FluxML/Flux.jl)
 - 動きました。ので、今回は Hello World してみましょう。 

# install

```
(v1.0 pkg> add Flux
```

基本的にはこれで済みますが、GPU使いたいときは

```
(v1.0 pkg> add CuArrays
```

をインストールしておきます。

# ドキュメントを読んでみます。

公式のドキュメント http://fluxml.ai/Flux.jl/stable/ を順に読みました。
サンプルもちょこちょこあるので、Chainer, PyTorch, TensorFlow, Keras などどれかのフレームワークを使っているような経験者であれば、雰囲気はつかめると思います。


# コードを書いてみやう

みんな大好きMNISTのExampleは

- https://github.com/FluxML/model-zoo/tree/master/vision/mnist

にて公開されています。ただ、関数化されてなかったり、学習済みモデルの保存、ロードまで書いてあるわけではないので、その他の機能は読者の演習問題のようです。

## 実装例

よなよな書いてこんな感じです。

```julia
#=
MNIST sample
Taken from model zoo of Julia and modified.
https://github.com/FluxML/model-zoo/blob/master/vision/mnist/mlp.jl
=#

using Printf

using Flux, Flux.Data.MNIST, Statistics
using Flux: onehotbatch, onecold, crossentropy, @epochs
using Base.Iterators: partition
using BSON: @load, @save
using CuArrays
using Random

function prepare_dataset(;train=true)
    train_or_test = ifelse(train,:train,:test)
    imgs = MNIST.images(train_or_test)
    X = hcat(float.(vec.(imgs))...)
    labels = MNIST.labels(train_or_test)
    Y = onehotbatch(labels,0:9)
    return X, Y
end


function define_model(;hidden)
    mlp = Chain(Dense(28^2,hidden,relu),
                Dense(hidden,hidden,relu),
                Dense(hidden,10),
                softmax)
    return mlp
end

function split_dataset_random(X, Y)
    divide_ratio=0.9
    shuffled_indices = shuffle(1:size(Y)[2])
    divide_idx = round(Int,0.9*length(shuffled_indices))
    train_indices = shuffled_indices[1:divide_idx]
    val_indices = shuffled_indices[divide_idx:end]
    train_X = X[:,train_indices]
    train_Y = Y[:,train_indices]
    val_X = X[:,val_indices]
    val_Y = Y[:,val_indices]
    return train_X, train_Y, val_X, val_Y
end

function train()
    println("Start to train")
    epochs = 10
    X, Y = prepare_dataset(train=true)
    train_X, train_Y, val_X,val_Y = split_dataset_random(X, Y)
    model = define_model(hidden=100) |> gpu
    loss(x,y)= crossentropy(model(x),y)
    accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))
    batchsize = 64
    train_dataset = gpu.([(train_X[:,batch] ,train_Y[:,batch]) for batch in partition(1:size(train_Y)[2],batchsize)])
    val_dataset = gpu.([(val_X[:,batch] ,val_Y[:,batch]) for batch in partition(1:size(val_Y)[2],batchsize)])

    callback_count = 0
    eval_callback = function callback()
        callback_count += 1
        if callback_count == length(train_dataset)
            println("action for each epoch")
            total_loss = 0
            total_acc = 0
            for (vx, vy) in val_dataset
                total_loss += loss(vx, vy)
                total_acc += accuracy(vx, vy)
            end
            total_loss /= length(val_dataset)
            total_acc /= length(val_dataset)
            @show total_loss, total_acc
            callback_count = 0
            pretrained = model |> cpu
            @save "pretrained.bson" pretrained
            callback_count = 0
        end
        if callback_count % 50 == 0
            progress = callback_count / length(train_dataset)
           @printf("%.3f\n", progress)
        end
    end
    optimizer = ADAM(params(model))

    @epochs epochs Flux.train!(loss, train_dataset, optimizer, cb = eval_callback)

    pretrained = model |> cpu
    weights = Tracker.data.(params(pretrained))
    @save "pretrained.bson" pretrained
    @save "weights.bson" weights
    println("Finished to train")
end

function predict()
    println("Start to evaluate testset")
    println("loading pretrained model")
    @load "pretrained.bson" pretrained
    model = pretrained |> gpu
    accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))
    println("prepare dataset")
    X, Y = prepare_dataset(train=false)
    X = X |> gpu
    Y = Y |> gpu
    @show accuracy(X, Y)
    println("Done")
end

function predict2()
    println("Start to evaluate testset")
    println("loading pretrained model")
    @load "weights.bson" weights
    model = define_model(hidden=100)
    Flux.loadparams!(model, weights)
    model = model |> gpu
    accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))
    println("prepare dataset")
    X, Y = prepare_dataset(train=false)
    X = X |> gpu
    Y = Y |> gpu
    @show accuracy(X, Y)
    println("Done")
end

function main()
    train()
    predict()
    predict2()
end

main()

```

# 何をやっているかの説明

## data loader
MNIST データセットをロードする関数はもう与えられているのでそれを使います。

```julia
julia> imgs = MNIST.images(:train)
julia> X = hcat(float.(vec.(imgs))...)
julia> labels = MNIST.labels(:train)
julia> Y = onehotbatch(labels,0:9)
julia> @show size(X)  size(Y);
size(X) = (784, 10000)
size(Y) = (10, 60000)
```

- 今回は単純なMLPを使うので２８ｘ２８の画像を一列のベクトルに直します。
Pythonと異なって `size(X)` の第一成分 は画像を構成する画素数28*28, 第二成分はデータの個数になっていることに注意します。画像処理もそうなのですが、Juliaでは列方向が優先されます。

- Y は ラベルをそのラベルの値が示す成分が１，それ以外は０とするOne-hot表現に変換しています。
ちなみにFluxでは onehotbatch の逆の操作をするonecoldbatch があります。（See: https://github.com/FluxML/Flux.jl/blob/master/docs/src/data/onehot.md)

あとは学習用とバリデーション用でデータをわけてバッチデータを作ります。

```julia
train_dataset = gpu.([(train_X[:,batch] ,train_Y[:,batch]) for batch in partition(1:size(train_Y)[2],batchsize)])
val_dataset = gpu.([(val_X[:,batch] ,val_Y[:,batch]) for batch in partition(1:size(val_Y)[2],batchsize)])
```

実際の運用ではバッチ学習するごとにデータを引っ張ってくるというのが多いですが、今回はデーターが軽いのでこんな書き方でもできちゃいます。

遅延して評価してくれるような設計を書きたい・・・

## model

単純なMLPモデルなので、下記のようにレイヤーをChainでつなげます。Chainer, KerasとかのSequentialみたいなものですね。

```julia
function define_model(;hidden)
    mlp = Chain(Dense(28^2,hidden,relu),
                Dense(hidden,hidden,relu),
                Dense(hidden,10),
                softmax)
    return mlp
end
```

loss や accuracy は次で定義しています。

```julia
loss(x,y)= crossentropy(model(x),y)
accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))
```

オプティマイザ `optimizer = ADAM(params(model))` も定義します。 

## 学習

```julia
epochs 10
@epochs epochs Flux.train!(loss, train_dataset, optimizer, cb = eval_callback)
```

`@epochs` マクロを用いて Flux.train!をepochs 分, ここでは１０, だけ呼び出します。
`train_dataset` の要素　`(train_X[:,batch] ,train_Y[:,batch])` 毎にパラメータを更新する
 のを１-イテレーションとする単位で学習を回しているんだと思います。`train_dataset` すべてまわったらそれを1-エポックとしてプログラム上では認識していると思われます。

ここは epoch を回すforとiterationをまわすforループを用いてバッチデータを入れた`Flux.train!`を書いても良いです。

```julia
#擬似コード
for e in 1:epochs
    shuffle!(dataset)
    for batch in partition(1:size(dataset)[2],batchsize)
        batch_data = dataset[batch]
        Flux.train!(loss, batch_data,opt,cb)
    end
end
```

`cb = eval_callback` は 1-iteration 毎に呼び出されるコールバック関数を引数にするようです。

```julia
    callback_count = 0
    eval_callback = function callback()
        callback_count += 1
        if callback_count == length(train_dataset)
            println("action for each epoch")
            total_loss = 0
            total_acc = 0
            for (vx, vy) in val_dataset
                total_loss += loss(vx, vy)
                total_acc += accuracy(vx, vy)
            end
            total_loss /= length(val_dataset)
            total_acc /= length(val_dataset)
            @show total_loss, total_acc
            callback_count = 0
            pretrained = model |> cpu
            @save "pretrained.bson" pretrained
            callback_count = 0
        end
        if callback_count % 50 == 0
            progress = callback_count / length(train_dataset)
           @printf("%.3f\n", progress)
        end
    end
```

## モデルの保存
一般的な運用を考えると学習したモデルは保存したいですよね。`BSON`というライブラリを用いてモデル、重みの保存、読み込みをするようです。

```julia
pretrained = model |> cpu
weights = Tracker.data.(params(pretrained))
@save "pretrained.bson" pretrained
@save "weights.bson" weights
println("Finished to train")
```

### |> gpu et |> cpu
説明してなかったですが、

`model = model |> gpu` で右辺の `model` をGPUに転送します。パイプ演算子 `|>` のかわりに`gpu(model)` と書いてもOKです。Chainerでいうところの `model.to_gpu()` みたいなものでしょうか。CPU側に戻したい場合は

`model = model |> cpu` と直します。モデルを保存するときは必ずCPUにしておかないとロードが失敗します。

## 推論

保存されたモデルを読み込む方法と、重みをロードして定義されたモデルに重みを注入する方法両方共できるようです。

```julia
@load "pretrained.bson" pretrained
```

たぶんですが、`@save` した時に用いた変数名を指定しないといけないようです。

重みだけ持っている場合は

```julia
@load "weights.bson" weights　# weights をロード
model = define_model(hidden=100)　#コードでモデルを定義
Flux.loadparams!(model, weights) #重みを注入
```

# 学習めんどくさい？

せやな、

- https://github.com/FluxML/Metalhead.jl　

とかあるで。

# 最後に

Flux.jl を深層学習用途で使う入門として MNIST を動かすことができました。学習の進め方、モデルの保存、ロード、推論のフローはつかめました。
Dense, CNN, DepthwiseConv, BatchNorm などのレイヤーが揃っているはずなので
他のモデルにも挑戦できそうです。

## 書いていてウッとなったこと
- `train` 関数内に (`model`, `X`, `Y` に依存するため) `loss`, `accuracy`, `callback`を書く設計になってしまっています。Julia力をつけて struct などを用いて改善したい。 
- 大規模なデータを扱う、バッチ学習をするために実画像のロードのルーチン、データを回す `Iterator` を実装しておかないと後々きつくなりそうです。 [ここで提案されています]((https://github.com/FluxML/Flux.jl/issues/150))。
- Lazy.jl とか使えるかな？