JAX+Flaxで始めるニューラルネットワーク

## 0. はじめに

JAXとは

- 自動微分
- GPU（TPU）
- Numpy互換性

な機能を持ったGoogle製の超便利なライブラリです。
詳しいところは、私も参考にした記事がいくつかあるので参照してください。

- [JAX入門～高速なNumPyとして使いこなすためのチュートリアル～](https://zenn.dev/koshian2/articles/af6758a5f3efc2)
- [JAXで始めるディープラーニング](https://note.com/moskomule/n/nc5fe3108ed0b)
- [JAX reference documentation](https://jax.readthedocs.io/en/latest/)

そして、
この記事ではJAXの上で動く、ニューラルネットワーク向けのライブラリ[Flax](https://flax.readthedocs.io/en/latest/)を使用したNNの実装方法を説明します。
あくまでも、NNを実装するための一例を紹介するという建前で、理論的な側面は割愛されています。
→ 動くNNモデルを実装できるようになることが、最終的な目標です。

https://jax.readthedocs.io/en/latest/

https://flax.readthedocs.io/en/latest/

## 1. 下準備

いきなり複雑なモデルは困るので、単純な回帰タスクを扱います。

### データセット

sklearnを使ってダウンロードできるcalifornia_housingを使います。
データの詳細は、[丁寧に解説してくれている記事](https://atmarkit.itmedia.co.jp/ait/articles/2201/31/news042.html)があるので見てみてください。

```python
import jax
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split

# READ
X, Y, *others = fetch_california_housing(as_frame=True).values()
feature_list = list(X.columns)

# SPLIT（70％を訓練、30％をテストデータとする）
X_TRAIN, X_TEST, Y_TRAIN, Y_TEST = train_test_split(X, Y, test_size=0.3, random_state=0)

# DEVICE_PUT
X_TRAIN, X_TEST, Y_TRAIN, Y_TEST = jax.device_put((
    X_TRAIN.to_numpy(),
    X_TEST.to_numpy(),
    Y_TRAIN.to_numpy(),
    Y_TEST.to_numpy(),
))

# SIZE
print("TRAIN:", X_TRAIN.shape, Y_TRAIN.shape)
print("TEST:", X_TEST.shape, Y_TEST.shape)
```

`jax.device_put`は（おそらく必須ではありませんが）GPU対応マシンではGPUデバイスにndarrayを転送してくれます。これによって、JAXで扱えるNumpyの互換型（＝DeviceArray）に変換されています。

### モデル

ここでは、超シンプルに、活性化関数にleaky_reluを使った全結合のNNを考えます。

<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2529987/59682df1-2156-5d76-9195-ec967ebfc171.png" width="75%">

ユニット数に注目すれば、8次元→8次元→4次元→2次元→1次元とマッピングをしていくだけの構造。
これをFlaxを使ったモデルに落とし込むと

```python
from typing import Sequence
from flax import linen as nn

class MLP(nn.Module):
    
    @nn.compact
    def __call__(self, X):

        for features in [8, 4, 2, 1]:
            X = nn.Dense(features=features)(X)
            X = nn.leaky_relu(X)

        return X.reshape(-1)
```

のようにクラスで定義することになります。
細かいところはFlaxのお作法として無理やり理解いただくとして、for文の中は

- nn.Dense: $N \times ?$次元の行列$X$をマッピングし、$N \times features$次元の行列を返す
- nn.leaky_relu: 行列の各セルに対して、leaky_reluを適用して返す ※Flaxにはその他の活性化関数も用意されています

が行われています。ここで$N$とはデータ数のことで、ミニバッチ学習ではバッチサイズになります。
つまり、今回のタスクでは

1. $N \times 8$次元の入力データを用意（＝入力のX）
1. $N \times 8$ にマッピング & 活性化関数を適用
1. $N \times 4$ にマッピング & 活性化関数を適用
1. $N \times 2$ にマッピング & 活性化関数を適用
1. $N \times 1$ にマッピング & 活性化関数を適用（→出力結果）

が行われているだけです。

注）最後に `.reshape(-1)` をしていますが、これは$N \times 1$の行列よりも$N$次元のベクトルのほうが（回帰では）使い勝手が良いというだけで、好みで変えてください。

### 損失関数

回帰タスクなのでMSEを採用します。

```python
import jax.numpy as jnp

def loss_function(params, X, Y):

    # 予測値
    pred = model.apply({'params': params}, X)

    # MSEを計算
    loss = jnp.mean((pred - Y)**2)

    return loss
```

（新しい変数名が出てきていますが、一旦スルーしてもらって大丈夫です。）

まず最初に、関数`model.apply`とモデルパラメータ`params`を使って、$N$個の入力データに関する予測値ベクトルを算出します。
そして、真値ベクトル$Y$を使って、MSEを計算し、損失`loss`を返します。

## 2. モデルパラメータの初期化

さて、前節でモデルを定義しましたが、パラメータをどのように管理していくか見ていきます。

まず、モデルを作成したいので、インスタンスを生成します。
```python
model = MLP()
```

そして、モデルの初期パラメータ`params`を
```python
key = jax.random.PRNGKey(0) # JAXではseedとしてkeyを発行します
params = model.init(key, X_TRAIN[:1])["params"]
```

のように得ます。`model.init`の第2引数は、訓練データの次元と合っていれば何でも良いので、とりあえず先頭行のみを入れています。
（正確に言うと、計算グラフを作るために仮のデータを流し込みます）

`params` はどのようになっているかというと、
```
FrozenDict({
    Dense_0: {
        kernel: DeviceArray([[-0.02331485, -0.64805007, -0.0834364 , -0.7591931 ,
                      -0.09689222, -0.546934  , -0.33758113, -0.5027879 ],
                     [ 0.19222587,  0.7038058 ,  0.11630198, -0.22404903,
                      -0.07720271,  0.56454706,  0.15410702,  0.43380636],
                     [ 0.7037437 , -0.7548236 , -0.20449154,  0.5431623 ,
                      -0.16591485, -0.6330585 , -0.1514996 , -0.30373544],
                     [ 0.19243844,  0.6275057 ,  0.38350856, -0.70182985,
                      -0.00386524,  0.41458794,  0.23042312, -0.6527162 ],
                     [ 0.32753906,  0.6719271 , -0.29333317, -0.28615567,
                       0.23924556,  0.09070686,  0.07164229,  0.6866846 ],
                     [-0.4248879 ,  0.04885302, -0.33519086,  0.23266272,
                       0.75276566,  0.20350276, -0.30964965,  0.07447256],
                     [ 0.20867735, -0.15052441, -0.14756508,  0.3662511 ,
                      -0.77800685, -0.62928206,  0.03433564,  0.12795487],
                     [ 0.36646786,  0.44371867, -0.30902764, -0.61118233,
                      -0.42007774,  0.15925379,  0.24388172,  0.11615389]],            dtype=float32),
        bias: DeviceArray([0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32),
    },
︙
    Dense_3: {
        kernel: DeviceArray([[-0.07865822],
                     [-0.00569075]], dtype=float32),
        bias: DeviceArray([0.], dtype=float32),
    },
})
```
のように、先ほどクラスの中で書いたnn.Denseに対応する、各層間の結合重み（kernel）とバイアス（bias）が辞書形式で入っています。
Flaxの面白いところは、モデル（のインスタンス）でパラメータを保持するのではなく、モデルの外で（変数`params`のように）管理します。

この初期パラメータを使って、実際に予測してみましょう。
インスタンス`model`の中にある関数`apply`を使って、`params`と入力データ行列を投げると、結果が返ってきます。
```python
pred = model.apply({"params": params}, X_TRAIN)
```
## 3. パラメータの更新

さて、SGDを考えると`params`を更新していくことになりますが、
$$ params = params - \epsilon \nabla E	 $$
を書いていくのは辛いです（さらにOptimizerにAdamを使ったりすると）。

そこで、パラメータ更新には、[optax](https://github.com/deepmind/optax)と呼ばれるライブラリと、Flaxの中にあるtrain_stateを用います。
今回はOptimizerにAdamを選び、
```python
import optax
from flax.training import train_state

# Adamを定義
tx = optax.adam(learning_rate=0.01)

# モデルパラメータの状態を定義
state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)

# 注）状態からパラメータを抽出するには params = state.params とする
```
としました。optaxでOptimizerを定義して、train_stateでパラメータを管理する「状態」を作成しています。

そして、損失関数`loss_function`を使って、
```python
# 勾配を計算（Xは入力データ行列，Yは真値）
grads = jax.grad(loss_function)(state.params, X, Y)

# パラメータ更新
state = state.apply_gradients(grads=grads)
```
のように、勾配を求めて、勾配を使ってパラメータの更新を行います。パラメータの更新式を直書きする必要なしに、`optax`と`train_state`を使うことでかんたんに書けます。

## 4. ミニバッチ学習

SGDでは、エポック毎に訓練データをランダムにミニバッチに分割していくのでした。
tensorflow-datasetsを使うなど色々なやり方がありますが、今回はJAXで完結させたいと思います。

先に完成形を提示すると
```python
@jax.jit
def train_batch(state, X, Y):

    # 勾配を計算
    grads = jax.grad(loss_function)(state.params, X, Y)
    # 更新
    state = state.apply_gradients(grads=grads)
    return state

epoch_nums = 64 # エポック数
batch_size = 512 # ミニバッチサイズ
batch_nums = X.shape[0] // batch_size # ミニバッチ数を計算

for epoch_idx in range(epoch_nums):

    key = jax.random.PRNGKey(epoch_idx)
    train_indices = jax.random.permutation(key, X_TRAIN.shape[0])

    for batch_idx in range(batch_nums):

        # ミニバッチ(X, Y)を抽出
        target_indices = jax.lax.dynamic_slice_in_dim(train_indices, (batch_idx*batch_size), batch_size)
        X, Y = X_TRAIN[target_indices], Y_TRAIN[target_indices]

        # ミニバッチで更新
        state = train_batch(state, X, Y)
```
のようになります。
まず、`train_batch`は、与えられたミニバッチ$X$, $Y$を用いてパラメータを更新し、更新された状態$state$を返す関数です。

そして、外のfor文がエポック単位の処理、中のfor文がバッチ単位の処理です。
まず、各エポックで訓練データのインデックスをシャッフル（＝`jax.random.permutation`）し、各バッチで、シャッフルされたインデックスを基に、部分データX, Yを抽出します。

## 5. おわり

ここまでに説明したものをキレイにしてまとめると、次のようになります。
```python
import jax
import jax.numpy as jnp
import optax
from flax import linen as nn
from flax.training import train_state
from typing import Sequence

from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split

"""
    モデル（クラス）の定義
"""
class MLP(nn.Module):
    
    @nn.compact
    def __call__(self, X):

        for features in [8, 4, 2, 1]:
            X = nn.Dense(features=features)(X)
            X = nn.leaky_relu(X)

        return X.reshape(-1)

"""
    損失関数
"""
@jax.jit
def loss_function(params, X, Y):

    # 予測値
    pred = model.apply({'params': params}, X)

    # MSEを計算
    loss = jnp.mean((pred - Y)**2)

    return loss

"""
    パラメータ更新（ミニバッチ単位）
"""
@jax.jit
def train_batch(state, X, Y):

    # 勾配を計算
    grads = jax.grad(loss_function)(state.params, X, Y)

    # 更新
    state = state.apply_gradients(grads=grads)
    return state

"""
    訓練データ・テストデータの用意
"""
# READ
X, Y, *others = fetch_california_housing(as_frame=True).values()
feature_list = list(X.columns)

# SPLIT（70％を訓練、30％をテストデータとする）
X_TRAIN, X_TEST, Y_TRAIN, Y_TEST = train_test_split(X, Y, test_size=0.3, random_state=0)

# DEVICE_PUT
X_TRAIN, X_TEST, Y_TRAIN, Y_TEST = jax.device_put((
    X_TRAIN.to_numpy(),
    X_TEST.to_numpy(),
    Y_TRAIN.to_numpy(),
    Y_TEST.to_numpy(),
))

"""
    メイン処理
"""
# モデルインスタンスの定義
model = MLP()

# 初期パラメータを生成
key = jax.random.PRNGKey(0)
params = model.init(key, X_TRAIN[:1])["params"]

# Adam Optimizer
tx = optax.adam(learning_rate=0.01)

# パラメータの状態を定義
state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)

# ハイパラの設定
epoch_nums = 64 # エポック数
batch_size = 512 # ミニバッチサイズ
batch_nums = X.shape[0] // batch_size # ミニバッチ数を計算

# バッチ学習
for epoch_idx in range(epoch_nums):

    key = jax.random.PRNGKey(epoch_idx)
    train_indices = jax.random.permutation(key, X_TRAIN.shape[0])

    for batch_idx in range(batch_nums):

        # ミニバッチ(X, Y)を抽出
        target_indices = jax.lax.dynamic_slice_in_dim(train_indices, (batch_idx*batch_size), batch_size)
        X, Y = X_TRAIN[target_indices], Y_TRAIN[target_indices]

        # ミニバッチで更新
        state = train_batch(state, X, Y)

    print(
        f"[{str(epoch_idx+1).zfill(2)}/{epoch_nums}]",
        "MSE（訓練）:", loss_function(state.params, X_TRAIN, Y_TRAIN),
        "MSE（汎化）:", loss_function(state.params, X_TEST, Y_TEST),
    )
```

以上です！お疲れ様でした！