この素晴らしいJuliaにDeep Learningを！

この記事はJulia Advent Calendar 2020 7日目の記事です。

イントロダクション: Julia言語とは何か?

• Pythonのように手軽に色んなスクリプトが書けて、
• Rみたいに統計解析が出来て、
• Perlのように文字列処理に強く、
• Lispのようにメタプログラミングをし、
• Matlabみたいに線形代数が扱えて、
• だけどC並みに高速であってほしいし、
• 分散・並列コンピューティングも出来てほしいし、
• REPLも欲しい、
• でもコンパイルもされてほしい…

そんな言語が存在するだろうか？　否、存在する。
Juliaはそんな「欲張りな人たち」のための言語である。

JuliaでDeep Learning

Juliaの勉強を兼ねて、書籍『ゼロから作るDeep Learning ❷ ―自然言語処理編』（オライリー・ジャパン）のPythonで実装されているコードの一部をJuliaに移植してみました。

現在、4章まで移植が完了しました。RNNの実装の直前ということでキリが良いので、1章の2層ニューラルネットワークの実装を紹介してみます。

その前に、JuliaとPythonの違いを少し見てみましょう。

PythonistaのためのJulia

PythonとJuliaで大きく異なる部分は何でしょうか？

実はJulia公式ドキュメントで特筆すべき点が紹介されているのですが、私がハマった点を少し紹介したいと思います。

クラスという概念は存在しない

Juliaは関数型言語であるため、そもそもオブジェクト指向プログラミング（OOP）という考え方をしていません。PythonでOOPに慣れきっていたため、当初は困惑していました。

しかし、クラスの代わりに構造体を定義し、クラスにメソッドを持たせるのではなく、構造体に関数を作用させていくように考えるとうまくいくことに気づきました。

Python版ではニューラルネットワークのレイヤーをクラスで定義していましたが、Juliaではそれを構造体（struct）の型（type）で定義します。そのため、レイヤーに持たせていたメソッドforward()とbackward()を独立させ、インスタンス変数paramsとgradsを構造体のインスタンス変数として持たせます。

独立させたメソッドはfunctionによって関数として定義します。ここで、「同じ名称の関数でどうやって処理を分岐させるんだ？」「if地獄になりそう…」と疑問が浮かぶかもしれません。Juliaには、型によって異なる関数を呼び出すことが出来る仕組みが標準で用意されています。これは、一般的に多重ディスパッチ（multiple dispatch）と呼ばれています。つまり、同じ名称の関数でレイヤーによって異なる処理を実行することが出来るのです。

配列は列指向（column major）

Python（NumPy）の配列は行指向（C由来）ですが、Juliaの配列はデフォルトでは列指向（Fortran由来）となっています。そのため、行指向の設計思想を基に実装されたアルゴリズムをそのまま列指向のシステムであるJuliaで実装すると思わぬ沼にハマります。

Python

>>> import numpy as np
>>> A = np.array([[1, 2, 3],[4, 5, 6]])
>>> for i in A:
...     print(i)
...
[1 2 3]
[4 5 6]

Julia

julia> A = [1 2 3; 4 5 6]
2×3 Array{Int64,2}:
 1  2  3
 4  5  6

julia> for i in A
           println(i)
       end
1
4
2
5
3
6

Aの列ごとに参照されていることが分かります。しかも、列ベクトルではなくスカラーが返ってきます。

Pythonのコードと同じ結果を得たい場合、次のように書きます。（あくまで一例です）

Julia

julia> for i in axes(A, 1)
           println(A[[i],:])
       end
[1 2 3]
[4 5 6]
julia> for i in axes(A, 1)
           println(typeof(A[[i],:]))
       end
Array{Int64,2}
Array{Int64,2}

似たような効果を得られる関数にBase.selectdimがあります。しかし、これを使用した場合、「$1×3$の行列」ではなく「3つの要素を持つ1次元配列」がSubArrayとして返ります。

selectdimを使用した場合

julia> for i in axes(A, 1)
           println(selectdim(A, 1, i))
       end
[1, 2, 3]
[4, 5, 6]
julia> for i in axes(A, 1)
           println(typeof(selectdim(A, 1, i)))
       end
SubArray{Int64,1,Array{Int64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true}
SubArray{Int64,1,Array{Int64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true}

型SubArray{}中の引数Int64,1から分かるように、数値型Int64の1次元配列となっています。

ちなみに、Python風のフォーマットで配列を記述すると、行列ではなく「配列の配列」となります。

Julia

julia> A = [[1, 2, 3], [4, 5, 6]]
2-element Array{Array{Int64,1},1}:
 [1, 2, 3]
 [4, 5, 6]

2層ニューラルネットワークによるデータセットの分離

線形分離不可能なデータセットとして有名なスパイラルデータセットを、2層ニューラルネットワークによって分離してみます。

1. スパイラルデータセットの生成

データセットは、sin関数とcos関数の和を取って生成します。

spiral.jl

module Spiral
using Random

export load_data

function load_data(;seed=1984)
    Random.seed!(seed)  # 関数直後の!は関数が破壊的変更を行うことを示す
    @show seed
    N = 100             # クラスごとのサンプルサイズ
    DIM = 2             # データの要素数
    CLS_DIM = 3         # クラス数

    x = zeros(N * CLS_DIM, DIM)
    t = zeros(Integer, N * CLS_DIM, CLS_DIM)

    for j = 0:CLS_DIM - 1 # Juliaではインデックスが1始まりのため
        for i = 0:N - 1
            rate = i / N
            radius = 1.0 * rate
            θ = j * 4.0 + 4.0 * rate + Random.randn() * 0.2

            ix = N * j + i
            x[ix + 1, :] = [radius * sin(θ) radius * cos(θ)]

            t[ix + 1, j + 1] = 1
        end
    end

    return x, t
end
end

GitHubリンク

生成したデータセットをGRバックエンドのPlotsで可視化したものはこちら（GitHub）。

2. 最適化アルゴリズムの実装

スパイラルデータセットの分離では、確率的勾配降下法（SGD）を使用します。

optimizer.jl

mutable struct SGD <: Optim
    """
    確率的勾配降下法
    Stocastic Gradient Descent
    """
    lr
    function SGD(;lr=0.01)
        new(lr)
    end
end

function update!(self::SGD, params, grads)
    for i = 1:lastindex(params)
        for j = 1:lastindex(params[i])
            params[i][j] = params[i][j] .- self.lr .* grads[i][j] # .-=は破壊的変更
        end
    end
end

GitHubリンク

コンストラクタメソッドで学習率lrの新しいオブジェクトを生成しています。SGDのアルゴリズムそのものはミニバッチをランダムに生成することに相当するので、この構造体ではオブジェクトの生成のみ行います。

3. 2層ニューラルネットワークの定義

まずはレイヤーの実装から。今回はAffineレイヤーと活性化関数のSigmoidレイヤーのみを使用します。

Affineレイヤーの実装です。引数layer::AffineでAffineレイヤーのforward関数とbackward関数を定義します。

layers.jl

mutable struct Affine <: AbstractLayer
    params
    grads
    x
    function Affine(W, b)
        layer = new()
        layer.params = [W, b]
        layer.grads = [zero(W), zero(b)]
        return layer
    end
end

function forward!(layer::Affine, x)
    W, b = layer.params
    layer.x = x
    return out = x * W .+ b
end

function backward!(layer::Affine; dout)
    W, b = layer.params
    dx = dout * W'
    dW = layer.x' * dout
    db = sum(dout, dims=1)

    layer.grads[1] = dW
    layer.grads[2] = db
    return dx
end

Sigmoidレイヤーの実装です。

layers.jl

mutable struct Sigmoid <: AbstractLayer
    params
    grads
    out
    function Sigmoid()
        layer = new()
        layer.params = []
        layer.grads = []
        return layer
    end
end

function forward!(layer::Sigmoid, x)
    return layer.out = 1 ./ (1 .+ exp.(-x))
end

function backward!(layer::Sigmoid; dout)
    return dx = dout .* (1.0 .- layer.out) .* layer.out
end

GitHubリンク

layer.jl内で実装したレイヤーを利用し、2層ニューラルネットワークを定義します。

two_layer_net.jl

# レイヤ定義ファイルの読み込み
include("../common/layers.jl")

using Random

# TwoLayerNetの定義
mutable struct TwoLayerNet
    layers
    loss_layer
    params
    grads
    
    function TwoLayerNet(;input_size, hidden_size, output_size, weight_init=0.01) # Keyword Argumentのみの場合はセミコロンが必要 cf. https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments
        self = new()
        I_s, H_s, O_s = input_size, hidden_size, output_size
        # 重みとバイアスの初期化
        W_1 = weight_init .* Random.randn(I_s, H_s)
        b_1 = zeros(H_s)' # JuliaのブロードキャストはNumPyと挙動が異なり，次元を追加しないため，予め転置して次元数を揃える
        W_2 = weight_init .* Random.randn(H_s, O_s)
        b_2 = zeros(O_s)'
        
        # レイヤの生成
        self.layers = [
            Affine(W_1, b_1),
            Sigmoid(),
            Affine(W_2, b_2)
        ]

        # すべての重みと勾配をまとめる
        self.params, self.grads = [], []
        for layer in self.layers
            push!(self.params, layer.params)
            push!(self.grads, layer.grads)
        end
        self.loss_layer = SoftmaxWithLoss()
        return self
    end
end

# Python版のTwoLayerNetのpredict()
function predict!(model::TwoLayerNet, x)
    for layer in model.layers
        x = forward!(layer, x) # Todo:同名モジュールをusingした場合でもモジュール名を省略したい
    end
    return x
end

# Python版のTwoLayersNetのforward()
function forward!(model::TwoLayerNet, x, t)
    score = predict!(model, x)
    loss = forward!(model.loss_layer, score, t) # SoftmaxWithLoss()のforward()
    return loss
end

# Python版のTwoLayersNetのbackward()
function backward!(model::TwoLayerNet; dout=1)
    dout = backward!(model.loss_layer, dout=dout) # SoftmaxWithLoss()のbackward()
    for layer in reverse(model.layers)
        dout = backward!(layer, dout=dout)
    end
    return dout
end

GitHubリンク

4. 学習

学習はJupyter上で実行しました。実際に学習を実行する関数のコードを抜粋してみます。

train_custom_loop.ipynb

function train!(model::TwoLayerNet, batch_x, batch_t)
    # 損失を計算
    loss = forward!(model, batch_x, batch_t) # model.forward()
    # 勾配を求める
    backward!(model) # model.backward()
    # パラメータを更新
    update!(optimizer, model.params, model.grads)
    
    push!(loss_list, loss)
    return loss
end

train_custom_loop.ipynb

function train!(model::TwoLayerNet, max_epoch, max_iters, train_x, train_t)
    total_loss = 0
    loss_count = 0
    
    for epoch = 1:max_epoch
        # ミニバッチの生成
        indices = rand(1:data_size, batch_size) # 1～data_sizeの範囲の一意な数値を要素に持つ要素数data_sizeのVectorを返す
        batch_x = train_x[indices, :]
        batch_t = train_t[indices, :]
        
        for iter = 1:max_iters
            loss = train!(model, batch_x, batch_t)
            total_loss += loss
            loss_count += 1
            # 定期的に学習経過を出力
            if iter % max_iters == 0
                avg_loss = total_loss / loss_count
                println("| epoch $(epoch) | iter $(iter) / $(max_iters) | loss $(round(avg_loss, digits=2))")
                push!(loss_list, avg_loss)
                total_loss, loss_count = 0, 0
            end
        end
    end
end

この上2つの関数を実行し、実際に学習を行うとPythonでは5.2秒のところ、3.8秒で学習が完了しました。

損失履歴をプロットしてみます。

少し怪しい感じもしますが、学習決定境界を描画してみます。

いい感じに分離出来ました！

このJupyter NotebookはGitHubで公開しています。また、全コードも私のGitHubリポジトリで公開しています。

まとめ

記事が冗長になってしまうのでDeep Learningと言いながら2層ニューラルネットワークの実装のみを紹介しましたが、Juliaの高速さの片鱗を知ってもらえたかと思います。イントロダクションで紹介した特長の他にも、推論のデプロイで利用できそうな実行可能なアプリケーションに変換してJuliaがインストールされていない環境でも動作させたり、もちろんGPUによる行列演算のアクセラレーションも可能な特長があります。機械学習フレームワークのFlux.jlなどもあります。

もちろんJuliaは機械学習のためのものだけではありません。確率的プログラミングもできるし、科学技術計算にもうってつけです。そのようなライブラリを追加するためのパッケージマネージャーも標準で用意されています。Pythonのようなクッソめんどくさい環境構築で悩む必要もないのです。ますますJulia界隈が活性化しそうですね。