Streamlitを使ってVAEをインタラクティブに表現する

概要

本記事は最近話題のStreamlitを使って、ディープラーニングの結果をインタラクティブに表現しようとした試みで、題材としてはVAEを用いています。

ディープラーニングの結果をweb上で公開しようと思ったら、本来ならバックエンドやJavaScriptを延々と書かいたうえに煩雑なデプロイもこなさなければなりませんが、Streamlitを使えばそれらが一瞬でできてしまいます。このStreamlitのすごさを体感してもらえればと思います。

VAEモデルの作成

題材として使うVAEモデルの作成を行います。【Python】Keras で VAE 入門 を参考にして作ってます。ここは本題ではないので、軽く流してもらって大丈夫です。

vae.py

import matplotlib.pyplot as plt
from tensorflow.keras.datasets import mnist
from keras.utils import np_utils
import numpy as np
import keras
from keras import layers
from keras import backend as K
from keras.models import Model
K.clear_session()

# MNISTを読み込み
(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train = x_train.astype('float32') / 255.
x_train = x_train.reshape(x_train.shape + (1,))
x_test = x_test.astype('float32') / 255.
x_test = x_test.reshape(x_test.shape + (1,))

# ラベルをカテゴリに
y_train = np_utils.to_categorical(y_train, 10)
y_test = np_utils.to_categorical(y_test, 10)

img_shape = (28, 28, 1)
img_shape_prod = np.prod(img_shape)
epochs = 30
batch_size = 256
latent_dim = 2

データセットはmnistを使っています。

モデルに関しては畳み込みを行うことにしました（パラメーターの調整が難しいので、精度だけ考えたら先例が豊富にある全結合層のみのほうがよかった泣）。

vae.py

# 潜在空間に落とし込む関数
def sampling(args):
    z_mean, z_log_var = args
    epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim), mean=0., stddev=1.)
    return z_mean + K.exp(z_log_var) * epsilon

# エンコーダー
encoder_input = layers.Input(shape=img_shape)
x = layers.Conv2D(32, 3, padding='same', activation='relu')(encoder_input)
x = layers.Conv2D(64, 3, padding='same', activation='relu', strides=(2, 2))(x)
x = layers.Conv2D(64, 3, padding='same', activation='relu')(x)
x = layers.Conv2D(64, 3, padding='same', activation='relu', strides=(2, 2))(x)
x = layers.Conv2D(64, 3, padding='same', activation='relu')(x)
conv_shape = K.int_shape(x)
x = layers.Flatten()(x)
x = layers.Dense(32, activation='relu')(x)

z_mean = layers.Dense(latent_dim)(x)
z_log_var = layers.Dense(latent_dim)(x)
z = layers.Lambda(sampling)([z_mean, z_log_var])

encoder = Model(encoder_input, [z_mean, z_log_var, z])
encoder.summary()

# デコーダー
decoder_input = layers.Input(shape=(latent_dim,))
x = layers.Dense(np.prod(conv_shape[1:]), activation='relu')(decoder_input)
x = layers.Reshape(conv_shape[1:])(x)
x = layers.Conv2DTranspose(64, 3, padding='same', activation='relu', strides=(2, 2))(x)
x = layers.Conv2DTranspose(32, 3, padding='same', activation='relu')(x)
x = layers.Conv2DTranspose(32, 3, padding='same', activation='relu', strides=(2, 2))(x)
x = layers.Conv2D(1, 3, padding='same', activation='sigmoid')(x)
decoder = Model(decoder_input, x)
decoder.summary()

# 結合
x = decoder(encoder(encoder_input)[2])
vae = Model(encoder_input, x)

VAE特有の処理として、z_mean, z_log_varを用いて、潜在空間（2次元）に落とし込んでいます。

vae.py

# 損失関数を定義
kl_loss = -0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
encoder_input = K.reshape(encoder_input, shape=(-1, img_shape_prod))
x = K.reshape(x, shape=(-1, img_shape_prod))
xent_loss = img_shape_prod * keras.metrics.binary_crossentropy(encoder_input, x)
vae_loss = K.mean(xent_loss + kl_loss)
vae.add_loss(vae_loss)

vae.compile(optimizer='rmsprop')
vae.summary()
history = vae.fit(x=x_train, y=x_train, shuffle=True, epochs=epochs, batch_size=batch_size, validation_data=(x_test, x_test))

# 保存
open("model.json", 'w').write(decoder.to_json())
decoder.save_weights('param.hdf5')

損失関数は関数を作成してvae.compile(optimizer='rmsprop', loss=loss_func)とするとうまくいかなかったですが、直接vae.add_loss(vae_loss)とするとうまくいきました。

vae.py

loss = history.history['loss']
val_loss = history.history['val_loss']

plt.plot(range(1,epochs), loss[1:], marker='.', label='loss')
plt.plot(range(1,epochs), val_loss[1:], marker='.', label='val_loss')
plt.legend(loc='best', fontsize=10)
plt.xlabel('epoch')
plt.ylabel('loss')
plt.show()
plt.savefig("loss")

from scipy.stats import norm

# 潜在空間上での出力画像の分布
n = 15
digit_size = 28
figure = np.zeros((digit_size * n, digit_size * n))
grid_x = norm.ppf(np.linspace(0.05, 0.95, n))
grid_y = norm.ppf(np.linspace(0.05, 0.95, n))

for i, yi in enumerate(grid_x):
    for j, xi in enumerate(grid_y):
        z_sample = np.array([[xi, yi]])
        z_sample = np.tile(z_sample, batch_size).reshape(batch_size, 2)
        x_decoded = decoder.predict(z_sample, batch_size=batch_size)
        digit = x_decoded[0].reshape(digit_size, digit_size)
        figure[i * digit_size: (i + 1) * digit_size,
               j * digit_size: (j + 1) * digit_size] = digit

plt.figure(figsize=(10, 10))
plt.imshow(figure, cmap='Greys_r')
plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
plt.axis("off")
plt.savefig("change_diagram")

モデルの訓練の結果はこんな感じです。

streamlitの導入

作成したモデルをstreamlitで書き下していきます。

st.py

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import streamlit as st
import os
from scipy.stats import norm
from keras.models import model_from_json

st.title("Streamlitを使ってVAEをインタラクティブに表示してみる")

"""
## 2次元の潜在空間上で変数を移動させることでインタラクティブに画像を生成する
"""

まず、タイトル（h1タグ）にするものをst.title(hogehoge)で書くことができます。
地の文は"""でコメントとして書くと反映されます。ここではマークダウンも使えます。

st.py

base_dir = os.path.dirname(os.path.abspath(__file__))

# サイドバー
st.sidebar.title('潜在変数を調整する')
xi = norm.ppf(st.sidebar.slider('x', 0.05, 0.95))
yi = norm.ppf(st.sidebar.slider('y', 0.05, 0.95))
batch_size = 256

st.sidebar.title("訓練時の潜在空間上での出力画像の分布")
img2 = Image.open(base_dir + "/change_diagram.png")
st.sidebar.image(img2, width=300)

折り畳み式のサイドバーも簡単に実装できます。基本的には要素を書く際にst.sidebar.としてあげるとサイドバーに表示されます。
画像についてはpillow形式ならst.image()で書くことができます。

st.py

# メイン
decoder = model_from_json(open(base_dir + "/model.json").read())
decoder.load_weights(base_dir + '/param.hdf5')
digit_size = 28

z_sample = np.array([[xi, yi]])
z_sample = np.tile(z_sample, batch_size).reshape(batch_size, 2)
x_decoded = decoder.predict(z_sample, batch_size=batch_size)
digit = x_decoded[0].reshape(digit_size, digit_size)
fig = plt.figure()
plt.axis("off")
plt.imshow(digit, cmap="Greys_r")
st.write(fig)

メインの部分はpillowで書くと画像として扱われてしまいぼやけた感じになってしまうので、ピクセルのような形で書けるmatplotlibで書くようにしています。

st.py

# チェックボックス
if st.checkbox('詳細情報を表示'):
	st.write(f"最大値　{digit.max():.4f}")
	st.write(f"最小値　{digit.min():.4f}")
	st.write(f"平均値　{digit.mean():.4f}")
	st.write(f"標準偏差{digit.std():.4f}")

最後に無駄にチェックボックスなんかを使ってみました。

デプロイ

デプロイはStreamlit sharingを使うことで簡単に行えますが、2つほど準備をする必要があります。
まず、requirements.txtを作成します。

requirements.txt

numpy==1.18.2
Pillow==7.1.1
matplotlib==3.2.1
Keras==2.4.3
scipy==1.4.1
streamlit==0.70.0
tensorflow-cpu==2.3.1

次にstreamlit sharingのページでデプロイの申請をする必要があります。申請が通るには2,3日かかります。(https://www.streamlit.io/sharing)

申請が通ったらデプロイすることができるようになります。
このリンク(https://share.streamlit.io/deploy )で、必要事項を記入してデプロイが完成です。

リンク

プロダクト：https://share.streamlit.io/eycjur/main/VAE/st.py

github：https://github.com/eycjur/main/tree/master/VAE

困った点

簡単に作れる分カスタマイズするのが難しいです。たとえばオブジェクトのマージンの設定やst.writeで大きさを調整する方法がわかりませんでした。私の理解力不足かもしれませんが、デザインを整えようとするとできない（or調べてもわからない）ことが多いので、既存のwebフレームワークとは使い分けをする必要があると思いました。

参考文献

• 【Python】Keras で VAE 入門