『ゼロから作る Deep Learning』を読んで Rust で実装した話

本記事は [Rust Advent Calendar 2020](https://qiita.com/advent-calendar/2020/rust) の６日目の記事です。
前日の記事は[cipepser](https://qiita.com/cipepser)様による[『Rustの可変長引数関数とHListの話』](https://cipepser.hatenablog.com/entry/rust-variadic-and-hlist)です。
翌日の記事は[termoshtt](https://qiita.com/termoshtt)様による[『Rust の Foreign Function Interface (FFI)』](https://qiita.com/termoshtt/items/0fa9959f9eb64b0907e2)です。

## 変更履歴
|変更ID|年月日|内容|
|:---:|:---|:---|
|1|2020/12/05|初版作成|
|2|2020/12/06|[namn1125](https://qiita.com/namn1125)様のコメントを受けて、比較結果を `cargo run --release` したものに修正|
|3|2020/12/06|レイヤーが受けるデータ型を Array2 に統一したときの訓練速度を比較結果に追記|

# 本記事のまとめ
* [『ゼロから作るDeep Learning―Pythonで学ぶディープラーニングの理論と実装』](https://www.oreilly.co.jp/books/9784873117584/)（以下、参考書と呼ぶことにします）を読み、5 章までの内容を [Rust](https://www.rust-lang.org/) で実装してみました。
* [MNIST](http://yann.lecun.com/exdb/mnist/) の手書き文字を認識する三層パーセプトロンモデルを実装する上で、筆者（私）が採用した方針と、ndarray を使う上で参考になりそうなポイントをリストアップします。

# リポジトリについて
* [こちらのリポジトリ](https://github.com/Surpris/rs-deep)に Rust で実装したコードが格納されております。
* [こちらの Oreilly Japan 社の GitHub リポジトリ](https://github.com/oreilly-japan/deep-learning-from-scratch)に元の Python コードが格納されています。
* [こちらのリポジトリ](https://github.com/Surpris/deep-learning-from-scratch)は上記 Oreilly Japan 社リポジトリを fork したものであり、ベンチマーク用に修正したコードと Keras in TensorFlow での実装が格納されております。

# 本編の目次
1. [動機](#動機)
1. [実装の基本方針](#実装の基本方針)
1. [モデルの実装](#モデルの実装)
    1. [データ型](#データ型)
    1. [レイヤーモデル](#レイヤーモデル)
    1. [順伝播と誤差逆伝播](#順伝播と誤差逆伝播)
1. [ポイント](#ポイント)
    1. [実装が必要な関数](#実装が必要な関数)
    1. [ArrayとArrayViewの変換](#arrayとarrayviewの変換)
    1. [Array2とArrayDの変換](#array2とarraydの変換)
    2. [Arrayから無作為に要素を抽出する方法](#arrayから無作為に要素を抽出する方法)
1. [ベンチマーク](#ベンチマーク)
1. [おわりに](#おわりに)
1. [参考](#参考)
    1. [先駆者の方々](#先駆者の方々)
    1. [Rustで使える深層学習フレームワーク](#rustで使える深層学習フレームワーク)

# 動機
これまで [Keras in TensorFlow](https://www.tensorflow.org/guide/keras/sequential_model) を使用して深層学習で遊ぶ側でしたが、内部の動作（例えば誤差逆伝播）の実装方法についてはからっきしでした。
ちょうど [Rust](https://www.rust-lang.org/) を勉強しているところであり、深層学習モデルの実装が Rust の学習において良い題材になると思い、実装にチャレンジしました。

# 実装の基本方針
今回の実装の基本方針は次の通りです。

1. できるだけ Rust-native なコードにする。   
    今回、他のプログラミング言語で作成されたライブラリに対する Rust bindings を作成しませんでした。また、Rust bindings の使用をできるだけ控えました。例えば [`NumPy`](https://numpy.org/) の Rust binding である [`rust-numpy`](https://github.com/PyO3/rust-numpy) は使用していません。
1. unsafe なコーディングは陽には行わない。   
    筆者がまだ Rust に慣れていないので、unsafe なコーディングを行わないこととしました。使用する既存のクレート内で unsafe なコーディングがされているものは許容しました。

# モデルの実装
## データ型
今回、特徴量とラベルのデータ型として [`ndarray`](https://github.com/rust-ndarray/ndarray) で実装されている型（`Array` 系）を採用しました。ndarray の利点として次が挙げられます。

* 多次元配列が実装されている。
* 代表的な配列演算が実装されている。

他の候補として配列と `Vec` が考えられましたが、以下に示す課題がありました。まず配列について、

* コンパイル時にサイズが確定している必要がある。
* 多次元配列の実装が面倒。
* 基本的な配列間の演算を自前で実装する必要がある。
    * [`numpy.dot`](https://numpy.org/doc/stable/reference/generated/numpy.dot.html) と同じ処理など。
* 関数が配列を戻り値にすることが難しい。
    * これは [`Box`](https://doc.rust-lang.org/std/boxed/struct.Box.html) でラップする、参照を引数にとる、等の方法で解決できると思われます。

次に Vec について、

* 多次元配列の実装が面倒。
* 基本的な配列間の演算を自前で実装する必要がある。
    * もしかしたら既存のクレートで実装があるかもしれません。
* 自前の実装では ndarray よりかなり処理速度が遅かった。
    * 実装次第では高速になると思います。

これらの課題を解決する手間と ndarray の利便性を鑑みて、ndarray を採用しました。

## レイヤーモデル
参考書の４章および５章で採用されているモデルはレイヤーモデルです。今回の実装もこれに従います。
今回実装した三層パーセプトロンモデルは下図のようなレイヤー構成を持ちます。
![三層パーセプトロンモデル](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/146354/46ab9e5e-651e-0852-3d32-c9d542cd2957.png)
&#8251;本記事において、$N$ 層のモデルとは、入力層、隠れ層、出力層の数の合計が $N$ であるモデルを指すものとします。

実装したレイヤー（構造体）を以下に挙げます。

* Affine 変換 ${\mathbf A}(\vec{x}) = {\mathbf W}\vec{x} + \vec{b}$
    * ${\mathbf W}$：重み、$\vec{b}$：バイアス
* 活性化関数 $\phi(x)$
    * ReLU：$\phi(x)=\\left\\{\\begin{array}{cc} x & (x > 0) \\\\  0 & (x \leq 0) \\end{array}\\right.$
    * Sigmoid：$\phi(x)=1/(1 + e^{-x})$
        * 今回のモデルでは使用していません。
    * Softmax：$\vec{x}=(x\_1,x\_2,\dots,x\_n)\rightarrow \phi(x_k\lvert\vec{x}) = e^{x\_k}/\sum\_j e^{x\_j}$
* 損失関数
    * Cross-entropy loss：$E(\vec{y}, \vec{t}) = -\sum\_k t\_k\log y\_k$
        * $\vec{y},\vec{t}$ はそれぞれ出力層の出力、ラベルを表します。
    * 簡単のため、softmax レイヤーと cross entropy error レイヤーを結合したレイヤーを採用しています。

各レイヤーに共通の関数はトレイトとして、次のような形式で実装しました。

```rust
pub trait LayerBase<T> {
    // 順伝播
    fn forward(&mut self, x: &ArrayD<T>) -> ArrayD<T>;
    // 逆伝播
    fn backward(&mut self, x: &ArrayD<T>, dx: &ArrayD<T>) -> ArrayD<T>;
    // レイヤーの重みの更新※
    fn update(&mut self, lr: T);
}
```

&#8251;今回実装した活性化関数レイヤーや損失関数レイヤーには必要ありませんが、便宜上共通関数として実装しました。

## 順伝播と誤差逆伝播
深層学習における出力の計算とモデルのパラメータの更新手順については、多くの方が記事にされています（Qiita 内であれば[こちらの記事](https://qiita.com/Ugo-Nama/items/04814a13c9ea84978a4c)や[こちらの記事](https://qiita.com/43x2/items/50b55623c890564f1893)が参考になると思います）。ここでは軽く触れるに留めます。

### 順伝播
$k$ 層目の隠れ層のサイズを $L\_k$ とします。次の二つの手順を繰り返すことで、入力 $\vec{y}\_1$ に対して予測値 $\vec{y}\_N$ を得ます。

1. Affine 変換レイヤー ${\mathbf A}\_k$ を通す：$\vec{v}\_{k} = {\mathbf A}\_k(\vec{y}\_{k}) = {\mathbf W}\_k\vec{y}\_{k} + \vec{b}\_k$
1. 活性化関数レイヤー ${\mathbf \Phi}\_k$ を通す：$\vec{y}\_{k+1} = {\mathbf \Phi}\_k(\vec{v}\_k)$

### 誤差逆伝播
パーセプトロンモデルで更新されるのは Affine変換レイヤーの重み ${\mathbf W}\_k = (w\_{ij}^{(k)})$ とバイアス $\vec{b}\_k = (b\_1^{(k)},b\_2^{(k)},\dots,b\_{L\_k}^{(k)})$ です。これらは次式のように損失関数の偏微分を用いて更新されます。
$$
w\_{ij}^{(k)}\rightarrow w\_{ij}^{(k)} - \epsilon\frac{\partial E}{\partial w\_{ij}^{(k)}},\ b\_j^{(k)} \rightarrow b\_j^{(k)} - \epsilon\frac{\partial E}{\partial b\_j^{(k)}}.
$$
ここで $\epsilon\ (> 0)$ は学習率です。
誤差逆伝播ですから、出力層に近い側から入力層に近い側へと誤差が伝播し、各層のパラメータが更新されます。
各層のパラメータの更新もまた層ごとに独立に実行できるため、update 関数を共通関数としてトレイトに実装しました。

さて、第 $N-1$ 層目に対応する Affine 変換レイヤーに対する誤差の偏微分を計算してみます。$E = E(\vec{y}\_N, \vec{t})$ は一見 ${\mathbf W}\_{N-1},\vec{b}\_{N-1}$ に依らないように見えますが、
$$
\vec{y}\_N = {\mathbf \Phi}\_{N-1}(\vec{v}\_{N-1}) = {\mathbf \Phi}\_{N-1}({\mathbf W}\_{N-1}\vec{y}\_{N-1} + \vec{b}\_{N-1})
$$
と表現できることを考慮すると、$E(\vec{y}\_N, \vec{t})$ が ${\mathbf W}\_{N-1},\vec{b}\_{N-1}$ に依存することが分かります。よって偏微分の連鎖律より
$$
\\begin{array}{rcl}
\frac{\partial E}{\partial w\_{ij}^{(N-1)}} &=& \frac{\partial E}{\partial \vec{y}\_N}\cdot\frac{\partial \vec{y}\_{N}}{\partial \vec{v}\_{N-1}}\cdot\frac{\partial \vec{v}\_{N-1}}{\partial w\_{ij}^{(N-1)}} = \sum\_l \sum\_m \frac{\partial E}{\partial y\_l^{(N)}}\frac{\partial y\_l^{(N)}}{\partial v\_m^{(N-1)}}\frac{\partial v\_m^{(N-1)}}{\partial w\_{ij}^{(N-1)}},\\\\
\frac{\partial E}{\partial b\_j^{(N-1)}} &=& \frac{\partial E}{\partial \vec{y}\_N}\cdot\frac{\partial \vec{y}\_{N}}{\partial \vec{v}\_{N-1}}\cdot\frac{\partial \vec{v}\_{N-1}}{\partial b\_j^{(N-1)}} = \sum\_l \sum\_m \frac{\partial E}{\partial y\_l^{(N)}}\frac{\partial y\_l^{(N)}}{\partial v\_m^{(N-1)}}\frac{\partial v\_m^{(N-1)}}{\partial b\_j^{(N-1)}}.
\\end{array}
$$
上式のうち $\partial E/\partial \vec{y}\_N$ は損失関数の偏微分、$\partial \vec{y}\_{N}/\partial \vec{v}\_{N-1}$ は活性化関数の偏微分、$\partial v\_m^{(N-1)}/\partial w\_{ij}^{(N-1)}$ および $\partial v\_m^{(N-1)}/\partial b\_j^{(N-1)}$ は Affine 変換の偏微分であり、それぞれ独立に計算できます。したがって、それぞれのレイヤーに偏微分操作を持たせるように誤差逆伝播の計算を実装することができます。このような背景があって、backward 関数を共通の関数としてトレイトに実装しました。

なお、残りのレイヤーのパラメータに対する誤差の偏微分も同様に計算できます。詳しくは[こちらの記事](https://qiita.com/Ugo-Nama/items/04814a13c9ea84978a4c)をご参照ください。

# ポイント
## 実装が必要な関数
便宜上、次の関数を `MathFunc` トレイトとして `ndarray::{Array1, Array2, ArrayD}` に実装しました。ArrayBase に実装するだけで良かったかもしれませんが、ちょっとうまく行きませんでした。この辺は将来的に修正したいところです。

* 指数関数：$x \rightarrow \exp(x)$
* 対数関数：$x \rightarrow \log_t(x)$
    * 自然対数を底とする対数関数：$x \rightarrow \log_e(x)$

```rust
pub trait MathFunc<T>{
    fn exp(&self) -> Self;
    fn log(&self, e: T) -> Self;
    fn log_natural(&self) -> Self;
}
```

今回、ラベルの形として one-hot 表現を採用しているため、予測結果を比較する際に `argmax` が必要になります。argmax は [`ndarray-stats`](https://docs.rs/ndarray-stats/0.3.0/ndarray_stats/) の [`QuantileExt`](https://docs.rs/ndarray-stats/0.3.0/ndarray_stats/trait.QuantileExt.html) トレイトに実装されているので、それを利用しました。

## ArrayとArrayViewの変換
今回、引数に ArrayD などの `Array<A, D> = ArrayBase<OwnedRepr<A>, D>` 系の型をとる関数を実装したのですが、ndarray の `slice` 系の関数が `ArrayView<'a, A, D> = ArrayBase<ViewRepr<&'a A>, D>` を返すので、その間の変換操作が必要になる場面が何度かありました。
ArrayView 系から Array 系に変換する関数が ndarray クレート内には無いようで、代替手段を見つけるのに苦労しました。
例えば次のように参照に定数を乗することで ArrayView を Array に変換できます。

```rust
let a: Array1<f32> = arr!(&[1.0, 2.0, 3.0]);
let b: Array1<f32> = 1.0 * &a.slice(s![..2]);
```

## Array2とArrayDの変換
今回、三層パーセプトロンモデルで２次元 Affine 変換を採用しました。`ndarray::Array2` には内積関数 `dot` が実装されていますが、ArrayD には実装されていません（もしかしたらテンソル積があるかもしれませんが、ちょっと見つけられていません）。
活性化関数などはできるだけ多次元配列に対応させたいということもあり、Array2 と ArrayD を互いに変換する必要が生じました。
ndarray クレート内に実装されている次の関数が前述の変換操作を実現します。

```rust
pub fn into_dimensionality<D2>(self) -> Result<ArrayBase<S, D2>, ShapeError>
where
    D2: Dimension, 
```

## Arrayから無作為に要素を抽出する方法
深層学習ではデータセットから無作為に抽出したサブセット（バッチ）を用いて訓練することがよくあります。Array に対して無作為抽出する方法を二つ紹介します。

### 方法１
一つ目は `rand` クレートを用いてインデックス配列を生成し、`Array::select` を用いる方法です。今回はこちらを使用しました。

```rust
// ArrayBase::select
pub fn select(&self, axis: Axis, indices: &[Ix]) -> Array<A, D>
where
    A: Copy,
    S: Data,
    D: RemoveAxis, 
```

例えば次のように用います。

```rust
use rand::prelude::*;

const BATCH_SIZE: usize = 100;
const NBR_TRAIN_IMAGES: usize = 60000;
fn main() {
    let data_set: Array2<f32> = ...;
    let mut rng = thread_rng();
    let mut indices: Vec<usize> = vec![0usize; BATCH_SIZE];
    for jj in 0..BATCH_SIZE {
        indices[jj] = rng.gen_range(0, NBR_TRAIN_IMAGES);
    }
    let batch: Array2<f32> = data_set.select(Axis(0), &indices);
}
```

### 方法２
二つ目は [`ndarray-rand`](https://github.com/rust-ndarray/ndarray/tree/master/ndarray-rand) の `sample_axis` を使う方法です。こちらは ndarray リポジトリ内で管理されているクレートです。

```rust
pub fn sample_axis(
    &self,
    axis: Axis,
    n_samples: usize,
    strategy: SamplingStrategy
) -> Array<A, D>
where
    A: Copy,
    D: RemoveAxis, 
```

こちらの使用例は [Docs.rs](https://docs.rs/ndarray-rand/0.12.0/ndarray_rand/trait.RandomExt.html#tymethod.sample_axis) や [ndarray-rand リポジトリの `tests/tests.rs`](https://github.com/rust-ndarray/ndarray/blob/master/ndarray-rand/tests/tests.rs) に掲載されています。

# ベンチマーク
Rust で三層パーセプトロンモデルを実装しましたが、元のコードでの実装と遜色ない性能を有してほしいところです。そこで以下のモデルを、訓練時間と確度（accuracy）の観点で比較しました。

* Rust でデータ構造に ndarray を採用して実装したもの
    * 本リポジトリの `ch05` に入っています。
    * `src/main.rs` で `rs_deep::dlfs01::ch05::train_neural_net::main();` 以外をコメントアウトした状態で、リポジトリのベースディレクトリ内で `cargo run --release` を実行すれば動きます。
* 元の Python コードでの実装
    * ベンチマーク用スクリプトは[こちら](https://github.com/Surpris/deep-learning-from-scratch/blob/master/ch05/train_neuralnet.py)です。
* Keras in TensorFlow を用いて実装したもの
    * ベンチマーク用スクリプトは[こちら](https://github.com/Surpris/deep-learning-from-scratch/blob/master/tensorflow/main.py)です。

## ベンチマークに用いた環境
|対象|値|
|:---|:---|
|CPU|Intel(R) Core(TM) i7-10510U @ 1.80 GHz|
|RAM|8.00 GB|
|OS|Ubuntu 20.04 on WSL2|
|元のコード用 Python|Python 3.7.9, NumPy 1.19.2|
|Keras in TensorFlow での実装用 Python|Python 3.7.9, NumPy 1.19.2, TensorFlow 2.3.0|
|Rust 用|rustup 1.23.0, rustc 1.48.0|

## パラメータ
|パラメータ|値|
|:---|:---|
|データセット|[MNIST](http://yann.lecun.com/exdb/mnist/)|
|入力層のサイズ $L\_1$|784 (=28 &times; 28)|
|隠れ層のサイズ $L\_2$|50|
|出力層のサイズ $L\_3$|10|
|学習率 $\epsilon$|0.1|
|訓練データセットの数|60,000|
|テストデータセットの数|10,000|
|バッチサイズ|100|
|イテレーション数|10,000|
|乱数の設定|なし|

## 結果
測定結果の一例を下表に示します。ここで
　訓練時間＝「バッチサイズ分のデータの選択＋順伝播＋逆伝播＋パラメータ更新」&times;「イテレーション数」
です。

|対象|CPU 使用率|訓練時間 (sec)|確度|
|:---|:---|:---|:---|
|元の Python コード|&sim;80%|&sim;10|train=0.9809, test=0.9708|
|Rust での実装（Array2 と ArrayD の混合）|&sim;20%|15.06 (debug だと 511)|train=0.9437, test=0.934|
|Rust での実装（すべて Array2）|&sim;20%|&sim;10|train=0.9479, test=0.9422|
|Keras in TensorFlow|&sim;70%|&sim;10|train=0.9813, test=0.9709|

Array2 と ArrayD の混合版では、訓練時間は約 1.5 倍かかりました。確度は0.035 程度劣ります。まだ実装に改善点がありそうです。
**（2020/12/06 追記）レイヤーが受けるデータ型をすべて Array2 に統一すると、訓練時間が Python での実装と同等になりました。**
また、CPU 使用率が Python のコードに比べて 1/4 位しかありません。CPU を有効に使用するという課題も見えてきました。

元の Python 実装が Keras での実装と同程度の訓練時間と精度が出ているのも興味深いです。
※Keras での実装は他の実装と完全に同一の操作が行われているわけではないので、その辺で多少の差がついているかもしれません。

# おわりに
今回は三層パーセプトロンモデルを Rust で実装し、その性能を元のコードや Keras in TensorFlow での実装と比較しました。
元のコードに比べて遅いなど、今回の実装には改善点が多数ありますので修正を進める予定です。
例えば Array2 と ArrayD の間の変換（`into_dimensionality`）がバッチごとに入るので、Array2 に統一することで速くなるかもしれません（フレキシビリティは下がりますが）。
**（2020/12/06 追記）実際、`into_dimensionality` に時間がかかるようです。**
本記事の執筆時点では５章（三層パーセプトロンモデルの実装）までしか進められておりませんが、参考書の残りの章も進めて基本的なレイヤーの実装方法も身に着ける予定です。

また、参考書には以下の続編があります。これらも読んで Rust で実装してみようと考えています。

* [『ゼロから作るDeep Learning―自然言語処理編』](https://www.oreilly.co.jp/books/9784873118369/)
* [『ゼロから作るDeep Learning―フレームワーク編』](https://www.oreilly.co.jp/books/9784873119069/)

# 参考
## 先駆者の方々
本記事と同等のことは誰もが考えるだろうと思っていましたが、やはり先駆者がいらっしゃいました。そのうち Qiita 内に投稿されている記事を紹介いたします。

* [[WIP]Rustで「ゼロからつくるDeep Learning」](https://qiita.com/eielh/items/5568d2f89924ce3adcef)（[eielh](https://qiita.com/eielh)氏）
* [ゼロから作るDeepLearning by Rust（第三章まで）](https://qiita.com/tkyk0317/items/27898219520e1d8cf6ff)（[tkyk0317](https://qiita.com/tkyk0317)氏）
* [RustでDeepLearning入門](https://qiita.com/ta_to_co/items/c5686bc702b9255b9c06)（[ta_to_co](https://qiita.com/ta_to_co)氏）
    * ArrayView の参照を引数にとる形式で関数が実装されています。

## Rustで使える深層学習フレームワーク
Rust で使える深層学習のフレームワークには、例えば以下のものがあります。

* [Leaf](https://github.com/autumnai/leaf)
    * 現在は開発が止まっています。
* [juice](https://github.com/spearow/juice)
    * 上記 Leaf の後継のようです。
* [primitiv-rust](https://github.com/primitiv/primitiv-rust)
    * [primitiv](https://github.com/primitiv/primitiv/) の Rust binding です。
    * [こちらの Qiita の記事](https://qiita.com/chantera/items/3efcd09256e6ec659620)にて紹介されています。
    * こちらも最近は開発が進んでいないようです。
* [Rust language bindings for TensorFlow](https://tensorflow.github.io/rust/tensorflow/index.html)
    * 公式の Rust binding です。
    * コアな部分は開発が進んでいるようですが、`examples` が充実していません。
    * また、examples を動かすためのモデルは Python でコンパイルされるようです。