Chainer のモデルを Flux.jl に変換するでキュ

# 本日は


- （プライベートなお話ですが）ベイズ勉強会お疲れ様でした．今日気づいたんですがChainer Meetupあったんですね・・・OTL
- 勉強会の後は飲み会でお肉を食べました．
- 裏mea/et up として Chainer のモデルを Juliaの深層学習フレームワークの Flux.jl の方に変換する方法を紹介します．

# Known Issues
- 2019/11/29日にFluxの 0.10 のバージョンアップで一部の機能が動作しないので Fluxのバージョンを0.9または0.8に下げてください．

# Gomah.jl を使う

Gomah.jl というJulia の野良パッケージがあります．ここでの野良はJuliaの公式パッケージとして登録されてないことを指します．
つまり

```julia
$ julia
julia> ] # パッケージモードに移行
(v1.1) pkg> add Gomah
```

では入らず次のようにする必要があります:

```julia
(v1.1) pkg> add https://github.com/terasakisatoshi/Gomah.jl.git
```

お気づきかと思いますが，私の自作パッケージです．

# Gomah.jl の使い方(導入）

- ChainerとJuliaの使い方に慣れている必要があります．
- 主な依存パッケージは PyCall.jl と Flux.jl です.

ひとまず使うには

1. Pythonの導入
1. Chainer,ChainerCVを導入. ターミナルを開いて`pip install numpy chainer chainercv`を実行するだけ
1. Juliaの導入 
1. JuliaのREPLを開く

```julia
$ julia
julia> ENV["PYTHON"]=Sys.which("python3")
julia> ] # パッケージモードに移行.
(v1.1) pkg> add PyCall Flux
(v1.1) pkg> add https://github.com/terasakisatoshi/Gomah.jl.git
(v1.1) pkg> Ctrl-C # REPLモードに移行
julia> using Gomah # Gomahパッケージをインポートする
[ Info: Recompiling stale cache file... (少し待つ)
julia> exit() # Juliaを終了
```

# モデルの変換

1. 検証用の画像として適当な画像，例えば，パイナップル画像を入手します．`pineapple.png` で作業用ディレクトリに保存
1. その作業用ディレクトリに移動して Julia を起動して下記のコードを実行（コピペでOK）
1. ChainerCV で用意されている訓練済みの `ResNet50` のモデルが Flux.jl の形式に変換されます．

```julia
using Gomah
using Gomah: L, np, reversedims
using Flux
using PyCall

using Test
using BenchmarkTools

py"""
import chainer
import chainercv
import numpy as np
num = 50
PyResNet = chainercv.links.model.resnet.ResNet
resnet = PyResNet(num, pretrained_model = "imagenet")
img=chainercv.utils.read_image("pineapple.png",dtype=np.float32,alpha="ignore")
img=chainercv.transforms.resize(img,(224,224))
_imagenet_mean = np.array(
            [123.15163084, 115.90288257, 103.0626238],
            dtype=np.float32
        )[:, np.newaxis, np.newaxis]
img=img-_imagenet_mean
img=np.expand_dims(img,axis=0)
resnet.pick=resnet.layer_names
with chainer.using_config('train', False):
    pyret=resnet(img)
    result=np.squeeze(pyret[-1].array)
    chidx=int(np.argmax(result))
    chprob=100*float(result[chidx])
    print(chidx)
    print(chprob)
"""

@testset "regression" begin
    num = 50
    myres = ResNet(num)
    Flux.testmode!.(myres.layers)
    img = reversedims(py"img")
    @show size(img), typeof(img)
    ret, name2data = myres(img)
    for (i,name) in enumerate(myres.layer_names)
        pyr = reversedims(py"pyret[$i-1].array")
        flr = name2data[name]
        @show name, size(flr)
        @test size(pyr) == size(flr)
        @show maximum(abs.(pyr .- flr))
    end
    flidx = argmax(ret)
    flprob = 100ret[argmax(ret)]
    @show flidx,flprob
    @test Int(py"chidx") == flidx[1]-1
    @show Float32(py"chprob") - flprob
end

@testset "benchmark" begin
    num=50
    img = reversedims(py"img")
    myres = ResNet(num)
    chainmodel = Chain(myres.layers...)
    Flux.testmode!(chainmodel)
    @time chainmodel(img)
    @time chainmodel(img)
    @time chainmodel(img)
    @time chainmodel(img)
    @time chainmodel(img)
    @time chainmodel(img)
end
```

実行結果で特にエラーがなければ成功です．Chainerのモデルを変換に成功できたようです．

# 何をしているか 

## JuliaのなかでPythonのコードを書く

- PyCallを用いると Julia から Python を呼ぶことができます. py"""Pythonのコード""" でJuliaのなかでPythonのスクリプトを記述できます．py"""Pythonのコード"""　の中にあるPythonのオブジェクト, 例えば `img` を得るには JuliaのREPLの上で 
 py"img" とすることでJuliaの（多次元）配列に変換されます．

## `ResNet(50)` を呼ぶ

`using Gomah` をすることで Goma.jl で定義されている ResNet というオブジェクトを利用することができます`ResNet(50)`でGomah.jlのなかでChainerCVのResNet(50) を変換するロジックが動きます．`Flux.testmode!.(myres.layers)` は Chainerでする `chainer.config.train=False` のFlux版だと思ってください．Fluxでは畳み込みネットワークに入力するデータフォーマットは `WHCN` です．これは Chainerのフォーマットの `NCHW` と全くの逆です．次元をひっくり返すために `reversedims` を使っています．あとは雰囲気で何をしているかわかると思います．ResNetのレイヤー（ResBlockとか）を通した時の出力をChainerと変換後の結果を比較しています．１０のマイナス５乗の誤差は出ますが分類問題を解く分には支障がないでしょう．

# 変換はどうしているの？

- 世の中AI,AIと騒がれていますが後ろでは人間が実装やマネージメントをしているようにChainerのレイヤーとFlux.jlの対応関係を人間（ここでは私が）眺めてChainerのオブジェクトから重みパラメータにアクセス-し対応するFluxのレイヤーに重みを移植します．結構原始的でしょ？

# 対応レイヤー

- ResNet を構成する基本的なレイヤーは対応させました．各々のフレームワークのコードをウーンと眺める＋入出力デバッグすればコツはつかめます．

## L.Linear -> Flux.Dense

ここでの `L` は Chainerユーザーにとって馴染みのある `import chainer.links as L` のことだと思ってください．
変換は下記のようにできます．

```julia
function ch2dense(link, σ = Flux.identity)
    W = link.W.array
    b = link.b.array
    Dense(W, b, σ)
end
```

`link` は L.Linear によって生成されたオブジェクトが入力されることを期待しています．

## L.Convolution2D -> Flux.Conv

下記のような変換をすれば良いです，

```julia
function ch2conv(link, σ = Flux.identity)
    # get weight W and bias b
    W = reversedims(link.W.array)
    # flip kernel data
    W = W[end:-1:1, end:-1:1, :, :]
    if py"hasattr"(link.b, :array)
        b = reversedims(link.b.array)
    else
        b = zeros(DTYPE, size(W)[4])
    end
    pad = link.pad
    s = link.stride
    Conv(W, b, σ, pad = pad, stride = s)
end
```

注意するべきことは，Chainerの重みとFluxの重みは配列の形状がちょうど逆さまになっていることです．例えば畳み込みのカーネルの情報を持つ `W` が Chainerの世界で `W.shape` が `(1,2,3,4)` だった場合，対応するFlux.Convの持つべき重みの形状 `size(W)` は `(4,3,2,1)` になっている必要があります．`nobias=True` なChainerオブジェクトの場合はバイアス項として形式的に零ベクトルを入れるようにします．同様に`DepthwiseConv` を用意していますので，気合いがあれば MobilenNet のレイヤーも構築することは原理的には可能となるでしょう．



# L.BatchNormalization -> Flux.BatchNorm

BatchNormalizationも可能です．Gomah.jl自体は学習済みモデルを変換しそれでJulia上で推論させることを目的としています．
ここでは `active` を `false` にしておいて作られたレイヤーはテストモードで動作するようにしています．
学習パラメータなどの変換は漏れがないように注意すれば対応関係は自然です．

```julia
function ch2bn(link, λ = Flux.identity)
    β = link.beta.array
    γ = link.gamma.array
    ϵ = Float32(link.eps)
    μ = link.avg_mean
    σ² = link.avg_var
    sz = size(μ)[1]
    momentum = Float32(0.1) #unused parameter

    #set `false` so as to BatchNorm works as test mode
    activate = false
    bn = BatchNorm(λ, β, γ, μ, σ², ϵ, momentum, false)
    return bn
end
```

あとは組み立てるだけですね．あとは　Max pooling などの変換に気をつけつつ下記のようにレイヤーを組みましょう．

```julia
struct ResNet
    layers
    layer_names
    function ResNet(num::Int)
        PyResNet = chainercv.links.model.resnet.ResNet
        pyresnet = PyResNet(num, pretrained_model = "imagenet")
        #dummyX = np.ones((1, 3, 224, 224), dtype = np.float32)
        #resnet(dummyX)
        #resnet101=ResNet(101, pretrained_model="imagenet")
        #resnet152=ResNet(152, pretrained_model="imagenet")
        @show pyresnet.layer_names
        layers = [
            Conv2DBNActiv(pyresnet.conv1),
            MaxPool((3, 3), pad = (0, 1, 0, 1), stride = (2, 2)),
            ResBlock(pyresnet.res2),
            ResBlock(pyresnet.res3),
            ResBlock(pyresnet.res4),
            ResBlock(pyresnet.res5),
            x -> reshape(mean(x, dims = (1, 2)), :, size(x, 4)),
            ch2dense(pyresnet.fc6),
            Flux.softmax,
        ]
        new(layers, pyresnet.layer_names)
    end
end

function (res::ResNet)(x)
    h = x
    d = Dict()
    for (name, lay) in zip(res.layer_names, res.layers)
        h = lay(h)
        d[name] = h
    end
    return h, d
end

```

`(res::ResNet)(x)` は　`ResNet` という自作の Julia の構造体に対して Chainerで言う所の `__call__(self,x)` に相当することが可能です．最近は `def forward(self, x)` とするのが良いみたいですね．

# まとめ

- Chainer で用意されている ResNet を変換して Julia 側で呼び出すことができたよ．
- 変換規則はコツを掴めば素直にできるよ．重みのフォーマットはChainerと相性がいいよ．

## この記事を読んで試してみたい人へ

- 推論に特化する場合は自動微分をサポートするような高級な Flux.jl に渡すよりももっと素朴なレイヤーを提供している NNlib.jl を使うといいかもしれません．Flux.jlの内部の実装を見ると NNlib.jl の関数に依存してるのがわかると思います．また, `@code_warntype` で調べるとLayerによっては BatchNorm の forward 実装に結構警告が出るのがわかると思います．ですので速度を追い求めてみたい人は Flux や NNlib.jl の実装を眺めて関数の引数に適切な型を与えることで Julia がJITしやすいようにする設計をすると良いでしょう．