Tensorflow 2.0.0rc0 + Tensorflow Probability 0.8.0rc0 -> RealNVP(1)

# Tensorflow 2.x + Tensorflow Probability
Tensorflow が 1.x から 2.x への ~~破壊的~~ 進行を開始し、世界中の人間がPyTorchへ移住しないかな、と願わずには居られない今日この頃ですね。

さて今回議題として扱うのは Tensorflow Probability というライブラリです。このライブラリは MCMC とか VAE とか GANs とか確率的な機械学習を行う際にその **確率的な** 部分をサポートするライブラリとして、Tensorflow から分離してできたライブラリです。

何が便利なのかっていわれると少々具体例に困りますが、ある確率分布をシミュレートしたり、ベイジアンなモデルを作る際なんかには使いやすくなるらしいです（制作陣談）。

問題となるのは、このライブラリ、Tensorflow 2.0 に対応したいらしく最新版の Tensorflow Probability 0.8.0rc0 から Tensorflow 2.0.0rc0 を requirements としているんですが、 **いくつかの関数が Tensorflow 2.0.0rc0 に対応していない** という欠陥を抱えています。

そのため、Tensorflow Probability を現状使うには、使える関数・クラスを継承して既存のクラスを参考にしながら **再実装** する必要があったりします。


# Flow base Model
Flow base Model は生成モデルの一種です。

生成モデルには、 VAE とか GANs とかあり、画像なんかでは GANs が猛威を奮っているイメージですが、 RealNVP はそこに一石を投じる足がかりになったものです。

Flow-base Model は データ $x$ を **可逆な関数 $f(x)$** を複数直列に適用して、元データに比べて解析しやすい分布 **z** に変換するという仕組みになっています。

$z$ が比較的操作しやすい分布(正規分布とか)になっているため、それを上手いこと操作したりすることで、良い感じの出力が得られます。

他の生成モデルが Flow-base Model に比べて抱えている点として、次の2つを挙げることができます。

- VAE は Reprezentation Trick などを使っているため、正確にデータから分布を移しているわけではない
- GANs は画像生成などの分野で高い性能を出しているが、定式化が難しい

軽く紹介した jupyter notebook は [ここ](https://github.com/MokkeMeguru/glow-realnvp-tutorial/blob/master/tips/flow-tutorial.ipynb) にあります。
# RealNVP

RealNVP は Flow base Model の一種です。 Google Brain から出されたモデルで、特徴は大きく分けて次の2つでしょうか。

1. (Affine) Coupling Layer
2. Multi-scale Architecture

(Affine) Coupling Layer は、特定軸方向にデータを2分割して、片方を片方でスケーリングする、みたいなことをします。定式化すると、次のようになります。

- $f(x)$

```math
\begin{eqnarray*}
    x_{1:d}, x_{d+1:D} &=& split(x) \\
    y_{1:d}&=&x_{1:d} \\
    y_{d+1:D}&=&s \odot x_{d+1:D} + t \\
    y &=& concat(y_{1:d}, y_{d+1:D}) \\
    where && \log{s}, t = NN(x_{1:d}) \\
\end{eqnarray*}
```

Multi-Scale Architecture は、変数変換（$x\leftrightarrow z$）を効率的に行うためのテクニックの一つで、一つ $f(x)$ が通ると出てきた $h$ を半分にして、半分を $z$ の一部として、もう片方を次の $f'(x)$ へ通すというものです。

簡単に図示すると、RealNVPの論文中の　Figure 4(b) がわかりやすいと思います。
 
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/147520/bf135348-a83f-f1f3-f4da-928c01968e37.png)


より理論的な話は、[notebook](https://github.com/MokkeMeguru/glow-realnvp-tutorial) を順番に見るか、 [RealNVPの論文を読んだので要約とメモ](https://qiita.com/exp/items/25e88f506fc7458c7fad) などを参照して下さい。

# Tensorflow で実装する
まず Single-Scale Architectureで RealNVP を実験してみましょう。

**tl;dr** [この notebook](https://github.com/MokkeMeguru/glow-realnvp-tutorial/blob/master/tips/RealNVP_tutorial.ipynb)

Notebookに必要な情報はすべて書いてある（はずな）のでこちらでは特に Coupling Layer の実装を例に、Tensorflow 2.0 と　Tensorflow Probability の使い方を紹介します。

```python

class RealNVP(tfb.Bijector):
    def __init__(
        self,
        input_shape,
        n_hidden=[512, 512],
        # this bijector do vector wise quantities.
        forward_min_event_ndims=1,
        validate_args: bool = False,
        name="real_nvp",
    ):
        """
        Args:
            input_shape: 
                input_shape, 
                ex. [28, 28, 3] (image) [2] (x-y vector)
                          
        """
        super(RealNVP, self).__init__(
            validate_args=validate_args, forward_min_event_ndims=forward_min_event_ndims, name=name
        )

        assert input_shape[-1] % 2 == 0
        self.input_shape = input_shape
        nn_layer = NN(input_shape[-1] // 2, n_hidden)
        nn_input_shape = input_shape.copy()
        nn_input_shape[-1] = input_shape[-1] // 2
        x = tf.keras.Input(nn_input_shape)
        log_s, t = nn_layer(x)
        self.nn = Model(x, [log_s, t], name="nn")

    def _forward(self, x):
        x_a, x_b = tf.split(x, 2, axis=-1)
        y_b = x_b
        log_s, t = self.nn(x_b)
        s = tf.exp(log_s)
        y_a = s * x_a + t
        y = tf.concat([y_a, y_b], axis=-1)
        return y

    def _inverse(self, y):
        y_a, y_b = tf.split(y, 2, axis=-1)
        x_b = y_b
        log_s, t = self.nn(y_b)
        s = tf.exp(log_s)
        x_a = (y_a - t) / s
        x = tf.concat([x_a, x_b], axis=-1)
        return x

    def _forward_log_det_jacobian(self, x):
        _, x_b = tf.split(x, 2, axis=-1)
        log_s, t = self.nn(x_b)
        return log_s
```

まず class ですが、Flow-base モデルの肝となる **可逆な関数** は `tfb.Bijector` を継承する必要があります。
もっと言ってしまうと、 `tfb.Bijector` に指定されている関数を実装すれば、可逆な関数を実装できる、というわけです。

## 引数

- input_shape    
 入力するデータの次元数です。例えば入力が 2次元ベクトルであるならば、 `input_shape=[2]` です。

- n_hidden    
 Coupling Layerでの、"半分の入力をもう半分でスケールさせる"、 際のスケールのために必要なNN層のための引数です。    
 尚NN層は主に2層の全結合層で構成されています。    

- forward_min_event_ndims
 これが Tensorflow Probability で重要となる引数です。これの仕様が複雑なので注意しなければなりません。
 forward_min_event_ndims は関数の適用軸数に合わせて指定する必要があります。例えば、element-wise に関数を適用する場合（例えば要素に対して正則化を施すとき）には 0, ベクトル毎に関数を適用する場合には 1 を指定します。

- validate_args
　基本的に False で大丈夫です。引数とかパラメータのバリデーションをしてくれるようですが、正直あんまり役立った場面に出会っていません。

- name
　その関数に名前を付けます。これはモデルに対して一意である必要があります。エラーが出た際にどの部分が悪いのかをすばやく見つけることができます。

## クラス継承 (#super)
super 関数は次のように定義されています。

```python
super(RealNVP, self).__init__(
            validate_args=validate_args, forward_min_event_ndims=forward_min_event_ndims, name=name
        )
```

引数が重要で、次のようなものがあります。
- validate_args: bool
  先述のとおりです。

- forward_min_event_ndims: int
  先述のとおりです。同様に inverse_min_event_ndims というものもありますが、特に両者が別でなければ、こちらだけ指定すれば良いです。

- name: str
  先述のとおりです。

- is_constant_jacobian: bool
  これは後述するヤコビアンが定数であるかどうかを指定します。こうすることで計算コストを下げることができます。Flow-base Modelは比較的にメモリをもりもり食うモデルなので、こういうテクニックは必須です。


## _forward 関数

```python
def _forward(self, x):
  x_a, x_b = tf.split(x, 2, axis=-1)
  y_b = x_b
  log_s, t = self.nn(x_b)
  s = tf.exp(log_s)
  y_a = s * x_a + t
  y = tf.concat([y_a, y_b], axis=-1)
  return y
```

これは $f(x)$ の関数を実装しています。わかりやすさのために、入力が $x$ 、出力が $y$ とすると良いでしょう。    
基本的には数式をそのまま書けば大丈夫です。

## _inverse 関数

_forward 関数と同様に作れば大丈夫です。但しオリジナル関数を作る際には、 **可逆になるように** 注意して実装して下さい。可逆さを確認するためには、クラスを実装してから次のようにして試すことができます。

```python
my_reversible_func = MyBijector(arg1, arg2, *args)
x = tf.random.normal(batch_shape)
y = my_reversible_func.forward(x)
rev_x = my_reversible_func.inverse(y)
print(tf.reduce_sum((rev_x - x) ** 2).numpy()) # near equal 0.0
```

## _forward_log_det_jacobian 関数

これは $log(det|\frac{d f(x)}{d x}|)$ 、つまり対数スケールのヤコビアンです。詳しい数式は [notebook](https://github.com/MokkeMeguru/glow-realnvp-tutorial/blob/master/tips/RealNVP_tutorial.ipynb) の方を参照して下さい。

_inverse_log_det_jacobian もありますが、これは一般に -1 * `forward log det jacobian` で示されることが多いので、省略することもできます。今回は省略しました。

# Tensorflow Probability における学習
Tensorflow Probability の、特に Flow base モデルの学習では、次のようにして学習を行うことができます。

## TransformedDistribution
```python
num_realnvp = 4
bijector_chain = []
for i in range(num_realnvp):
    bijector_chain.append(RealNVP(input_shape=[2], n_hidden=[256, 256]))
    bijector_chain.append(tfp.bijectors.Permute([1, 0]))

flow = tfd.TransformedDistribution(
    distribution=mvn,
    bijector=tfb.Chain(list(reversed(bijector_chain)))
)
print('trainable_variables: ', len(flow.bijector.trainable_variables))
```

`TransformedDistribution` はある分布をある分布へ写したい、つまり Flow-base Model のようなモデルのためのクラスです。
引数として、

- distribution
  変換後の分布、例えば正規分布とか

- bijector
  $f(x)$

- (event_shape　次回登場します)
  データの次元数

などがあります。

## 損失関数
```python
@tf.function
def loss(targets):
    return - tf.reduce_mean(flow.log_prob(targets))


optimizer = tf.optimizers.Adam(learning_rate=0.001) 
log = tf.summary.create_file_writer('checkpoints')
avg_loss = tf.keras.metrics.Mean(name='loss', dtype=tf.float32)
```

`tf.function` デコレータは Python の関数を tensorflow で効率化された形にコンパイルするためのデコレータです。これによって学習速度が 10 倍くらい(体感。でも多分この程度)になるので **本番** ではつけるようにしましょう。

但しはじめてモデルにデータを通す際には、このデコレータは外して下さい。 Tensorflow で読める形式にコンパイルされる都合上、 **エラーが起きたときのメッセージが摩訶不思議なことになる** ためです。

損失関数自体はお好みのものを指定すれば良いと思いますが、おそらくこれが最もスタンダードです。

## 学習
```python

n_epochs = 200

for epoch in range(n_epochs):
    for targets in train_dataset:
        with tf.GradientTape() as tape:
            log_prob_loss = loss(targets)
        grads = tape.gradient(log_prob_loss, flow.trainable_variables)
        optimizer.apply_gradients(zip(grads, flow.trainable_variables))
        avg_loss.update_state(log_prob_loss)
        if tf.equal(optimizer.iterations % 1000, 0):
            with log.as_default():
                tf.summary.scalar("loss", avg_loss.result(), step=optimizer.iterations)
                print(
                    "Step {} Loss {:.6f}".format(
                        optimizer.iterations, avg_loss.result()
                    )
                )
                avg_loss.reset_states()
```

今の所自己流ですが、こんな感じでやると良いと思います。(公式のTutorialがそのうち[出てくるらしい](https://github.com/tensorflow/probability/issues/554)) ので、それに従って下さい。）

簡単に言うと、データソースからバッチを取り出して(2つ目のfor iteration) 、 `gradientTape` を貼って損失を計算して、勾配から最適化関数を訓練可能な変数に適用します。

if 文は学習ログを出力するための部分で、ここを上手いこと書き換えることで、Tensorboard で訓練を可視化することができます。(次回以降では実装しています。)

## 結果
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/147520/e8a98dd0-8fc9-fdcb-1aac-3d264fbaa0ff.png)

可逆な変換がうまく行っていますね。

# TODO
今の所、RealNVPとGlowしか実装していないので、
MAFとかの実装もしてみたいです。

~~あと実装が合っているか不安~~

# 次の話
[Tensorflow 2.0.0rc0 + Tensorflow Probability 0.8.0rc0 -> RealNVP(2)]
(https://qiita.com/MeguruMokke/items/6d3beddce93af6cd7f66)
[Tensorflow 2.0.0rc0 + Tensorflow Probability 0.8.0rc0 -> Glow](https://qiita.com/MeguruMokke/items/d67d78d41977c4d22541)