Julia Fluxでmnist

## Flux

[Flux](https://fluxml.ai/Flux.jl/stable/)は、Julia用の機械学習ライブラリ。自動微分とかも入っていていい感じにNNが使える、らしい。試しにMNISTをやってみる。

## MNISTの読み込み
[MNIST](http://yann.lecun.com/exdb/mnist/)は言わずとしれた手書き文字データセット。色々なフォーマットであちこちで配布されているが、ここでは、このサイトからバイナリをダウンロードして使ってみる。フォーマットは、ページの下の方に書かれているが、32bit整数がMSB firstで書かれているので、MSB lastのIntelチップでは入れ替えなくちゃいけなくて面倒。きっと、SPARCとか使ってたんだろうなあ。

```Julia
using GZip

prefix = "mnist/"
test_x_f = "t10k-images-idx3-ubyte.gz"
test_y_f = "t10k-labels-idx1-ubyte.gz"
train_x_f = "train-images-idx3-ubyte.gz"
train_y_f = "train-labels-idx1-ubyte.gz"

## バイト順の入れ替え
function reverseMSB(val::Int32)
    (val & 0xff000000) >> 24 +
    (val & 0x00ff0000) >> 8  +
    (val & 0x0000ff00) << 8  +
    (val & 0x000000ff) << 24 
end

## ラベルファイルの読み込み
function read_label(filename::String) 
    GZip.open(filename) do f
        _ = read(f, Int32);
        len = reverseMSB(read(f, Int32));
        read(f)
    end
end

## データファイルの読み込み
function read_data(filename::String) 
    GZip.open(filename) do f
        _ = read(f, Int32);
        len = reverseMSB(read(f, Int32));
        rows = reverseMSB(read(f, Int32));
        cols = reverseMSB(read(f, Int32));
        bytes = read(f);
        reshape(bytes, (rows, cols, len))
    end
end

##
train_x = read_data(prefix * train_x_f)
train_y = read_label(prefix * train_y_f)
test_x = read_data(prefix * test_x_f)
test_y = read_label(prefix * test_y_f)

## データファイルを整形。バイト値から浮動小数点に変換
numdata = size(train_x)[3]
train_x0 = [reshape(train_x[:, :, i], :) for i in 1:numdata] ./ 255.0
test_x0 = [reshape(test_x[:, :, i], :) for i in 1:size(test_x)[3]] ./ 255.0
```

## バッチの作成
Fluxの`train!`は、バッチのイテレータを引数として受け取る。なので、バッチを生成するイテレータを書いてやる必要がある。サイズ小さいので、単純にバッチの配列を渡すだけでもいいような気もするけど、一応イテレータで書いてみた。

``` Julia
mutable struct Batch
    x::Vector
    y::Vector
    size::Int64
    batch_size::Int64
    Batch(a, b, size) =
        new(a, b, length(a), size)
end

function Base.iterate(self::Batch, state=1)
    if state == 1 
        indices = randperm(self.size)
        self.x = self.x[indices]
        self.y = self.y[indices]
    end
    if state + self.batch_size - 1 > self.size 
        return nothing
    end
    b_end = state+self.batch_size - 1
    next_batch = (hcat(self.x[state:b_end]...), Flux.onehotbatch(self.y[state:b_end], 0:9))
    return (
        next_batch,
        state + self.batch_size
    )
end
Base.length(self::Batch) = Integer(floor(self.size / self.batch_size))
Base.eltype(self::Batch) = Tuple
```

## モデル

モデルは`Flux.Chain`でネットワークを積層する。構造体にする意味あんまり無いような気もするが、とりあえず、構造体にまとめてみた。実際のネットワーク構成は、構造体作成時に引数で与える。lossはクロスエントロピー。

``` julia
import Flux
import Flux.params
struct NN
    layers
    NN(array) = begin
        l = Flux.Chain(array...)
        new(l)
    end
end    
params(self::NN) = Flux.params(self.layers)
loss(self::NN, x, y) = Flux.crossentropy(self.layers(x), y)
```

## 訓練と評価
訓練は、バッチを渡して`train!`を呼び出すだけ。その際に、lossを計算する関数と、パラメータを返す関数を指定する。便利な世の中になったものだ。

``` julia
function train(nn, train_t, test_t, num_epochs=10)
    for epoch in 1:num_epochs
        Flux.train!(
            (x, y)-> loss(nn, x, y), 
            params(nn),  
            Batch(train_t[1], train_t[2], 100),
            Descent())
        println("epoch=", epoch, ": ", compute_loss_acc(nn, test_t[1], test_t[2]))
    end
end
```

`compute_loss_acc`はロスの平均値と予測精度を返す。上のtrainの中から、1エポックに1度呼び出される。

``` julia
function compute_loss_acc(nn, x, y) 
    loss_sum = 0.0;
    hit_sum = 0;
    total = 0;
    for (x0, y0) in Batch(x, y, 100)
        y_bar = nn.layers(x0)
        loss_sum += Flux.crossentropy(y_bar, y0)
        #println(y_bar)
        hit_sum += sum((getindex.(argmax(y_bar, dims=1), 1))[1,:] .== Flux.onecold(y0))
        total += size(y_bar)[2]
    end
    loss_sum/total, hit_sum/total
end 
```

## 実行
まずはモデルを作る。面倒なので、コンボリューションではなく全結合層を2つ重ねる。活性化関数はRELU。最後にsoftmax。

``` julia
nn = NN([
    Dense(28*28, 30, Flux.relu), 
    Dense(30, 10, Flux.relu), 
    Flux.softmax
])
```

で、訓練データとテストデータを引数に呼び出す。

``` julia
train(nn, (train_x0, train_y), (test_x0, test_y), 10)
```

結果はこんな。

``` 
epoch=1: (0.002921946030484957, 0.9177)
epoch=2: (0.0024627857710172777, 0.9282)
epoch=3: (0.0021479621613419773, 0.9373)
epoch=4: (0.0019236276361788462, 0.9442)
epoch=5: (0.0017464351921651347, 0.9474)
epoch=6: (0.0016167659175639116, 0.9525)
epoch=7: (0.0015334977051955692, 0.9541)
epoch=8: (0.0014695846478115899, 0.9557)
epoch=9: (0.0014405294374969476, 0.9571)
epoch=10: (0.0014510845261526148, 0.9575)
```

## まとめ
GPUないので10エポックで30秒弱ぐらい。こんな雑なモデルでも精度95%ぐらいは出るのは驚き。

機能的にはPyTorchやらTensorFlowやらとは比べられないが、とりあえず一通りの道具は揃っていそう。