[Python][Keras]VAEの潜在変数空間結果にうまく違いが現れません。何が原因なのかご教授お願い致します。

Question

前提

初学者なので至らない点があると思います、申し訳ございません。

PythonでKerasを用いて機械学習(VAE(Variational AutoEncoder))をGooglecolab上で行いました。
Mnistデータセットではなく、オリジナル画像(1100枚)でVAEを実行し、潜在変数空間を可視化したところ、潜在変数内で出力画像にあまり変化が見られませんでした。

この原因が分からず、質問させていただきました。

原因として私が考えたもの
①Mnistデータセットと比べるとサンプル数が少ない
②サンプル画像に見た目の差異が少ない
③(256×256ピクセルのとき)CNNの構造が上手くいっていない
のどれかであると考えました。
ご教授いただければ幸いです。

実現したいこと

上手く潜在変数空間で変化している出力結果を得たい。

発生している問題・エラーメッセージ

該当のソースコード

#①(Mnistと同様に28×28ピクセルでサンプル画像を加工したもの)
from tensorflow import keras
from tensorflow.keras import layers

latent_dim = 2

encoder_inputs = keras.Input(shape=(28, 28, 1))
x = layers.Conv2D(32, 3, activation="relu", strides=2, padding="same")(encoder_inputs)
x = layers.Conv2D(64, 3, activation="relu", strides=2, padding="same")(x)
x = layers.Flatten()(x)
x = layers.Dense(16, activation="relu")(x)
z_mean = layers.Dense(latent_dim, name="z_mean")(x)
z_log_var = layers.Dense(latent_dim, name="z_log_var")(x)
encoder = keras.Model(encoder_inputs, [z_mean, z_log_var], name="encoder")

encoder.summary()

import tensorflow as tf
class Sampler(layers.Layer):
    def call(self, z_mean, z_log_var):
        batch_size = tf.shape(z_mean)[0]
        z_size = tf.shape(z_mean)[1]
        epsilon = tf.random.normal(shape=(batch_size, z_size))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

latent_inputs = keras.Input(shape=(latent_dim,))
x = layers.Dense(7 * 7 * 64, activation="relu")(latent_inputs)
x = layers.Reshape((7, 7, 64))(x)
x = layers.Conv2DTranspose(64, 3, activation="relu", strides=2, padding="same")(x)
x = layers.Conv2DTranspose(32, 3, activation="relu", strides=2, padding="same")(x)
decoder_outputs = layers.Conv2D(1, 3, activation="sigmoid", padding="same")(x)
decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

decoder.summary()

class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.sampler = Sampler()
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = keras.metrics.Mean(
            name="reconstruction_loss")
        self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")

    @property
    def metrics(self):
        return [self.total_loss_tracker,
                self.reconstruction_loss_tracker,
                self.kl_loss_tracker]

    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var = self.encoder(data)
            z = self.sampler(z_mean, z_log_var)
            reconstruction = decoder(z)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(
                    keras.losses.binary_crossentropy(data, reconstruction),
                    axis=(1, 2)
                )
            )
            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            total_loss = reconstruction_loss + tf.reduce_mean(kl_loss)
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        return {
            "total_loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
        }
#ファイルからサンプル画像を取り出し、28×28に加工
from IPython.display import Image, display
import glob
import os
import re
import cv2
import numpy as np
x=glob.glob('{}/*.png'.format("2dplot_Wing_ex_1100"))
image_28 = []
for i in x:
  img = cv2.imread(i)
  img_gray= cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  img_gray_28=cv2.resize(img_gray,(28,28))
  img_gray_asarray_28 = np.asarray(img_gray_28)
  image_28.append(img_gray_asarray_28)

image_np_array_2828=np.array(image_28)
(x_train), (x_test) = image_np_array_2828[:-100],image_np_array_2828[-100:]#1100枚の内train1000枚、test100枚
x_train = x_train.astype('float32') / 255.
x_train = np.reshape(x_train, (len(x_train), 28, 28))
x_test = x_test.astype('float32') / 255.
x_test = np.reshape(x_test, (len(x_test), 28, 28))

import numpy as np#学習
mnist_digits = np.concatenate([x_train, x_test], axis=0)
mnist_digits = np.expand_dims(mnist_digits, -1).astype("float32") / 255
vae = VAE(encoder, decoder)
vae.compile(optimizer=keras.optimizers.Adam(), run_eagerly=True)
vae.fit(mnist_digits, epochs=30, batch_size=10)

import matplotlib.pyplot as plt

n = 20
digit_size = 28
figure = np.zeros((digit_size * n, digit_size * n))

grid_x = np.linspace(-1, 1, n)
grid_y = np.linspace(-1, 1, n)[::-1]

for i, yi in enumerate(grid_y):
    for j, xi in enumerate(grid_x):
        z_sample = np.array([[xi, yi]])
        x_decoded = vae.decoder.predict(z_sample)
        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=(30, 30))
start_range = digit_size // 2
end_range = n * digit_size + start_range
pixel_range = np.arange(start_range, end_range, digit_size)
sample_range_x = np.round(grid_x, 1)
sample_range_y = np.round(grid_y, 1)
plt.axis("off")
plt.imshow(figure, cmap="Greys_r")

#②(256×256ピクセルでサンプル画像を加工)
#エンコーダ部分
from tensorflow import keras
from tensorflow.keras import layers

latent_dim = 2

encoder_inputs = keras.Input(shape=(256, 256, 1))
x = layers.Conv2D(32, 3, activation="relu", strides=2, padding="same")(encoder_inputs)
x = layers.Conv2D(64, 3, activation="relu", strides=2, padding="same")(x)
x = layers.Flatten()(x)
x = layers.Dense(16, activation="relu")(x)
z_mean = layers.Dense(latent_dim, name="z_mean")(x)
z_log_var = layers.Dense(latent_dim, name="z_log_var")(x)
encoder = keras.Model(encoder_inputs, [z_mean, z_log_var], name="encoder")
#デコーダ部分
latent_inputs = keras.Input(shape=(latent_dim,))
x = layers.Dense(64 * 64 * 64, activation="relu")(latent_inputs)
x = layers.Reshape((64, 64, 64))(x)
x = layers.Conv2DTranspose(64, 3, activation="relu", strides=2, padding="same")(x)
x = layers.Conv2DTranspose(32, 3, activation="relu", strides=2, padding="same")(x)
decoder_outputs = layers.Conv2D(1, 3, activation="sigmoid", padding="same")(x)
decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")
#画像を加工し、代入する部分
from IPython.display import Image, display
import glob
import os
import re
import cv2
import numpy as np

x=glob.glob('{}/*.png'.format("2dplot_Wing_ex_1100"))
image_256 = []
for i in x:
  img = cv2.imread(i)
  img_gray= cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  img_gray_256=cv2.resize(img_gray,(256,256))
  img_gray_asarray_256 = np.asarray(img_gray_256)
  image_256.append(img_gray_asarray_256)
image_np_array_256256=np.array(image_256)
(x_train), (x_test) = image_np_array_256256[:-100],image_np_array_256256[-100:]
x_train = x_train.astype('float32') / 255.
x_train = np.reshape(x_train, (len(x_train), 256, 256))
#print(x_train.shape) →(300, 256, 256)
x_test = x_test.astype('float32') / 255.
x_test = np.reshape(x_test, (len(x_test), 256, 256))
#潜在変数のdigit_size=256に変更

試したこと

・epochやbatchsizeは色々変えて試しましたが、あまり変化は得られなかった。
・256×256ピクセルのデコーダ、エンコーダ層の構造があやしい？あまり詳しくないため、すみません

補足情報（FW/ツールのバージョンなど）

・各サンプル画像は下記4つの図形を係数を少しずつ変化させてサンプル画像(各350枚用意した(matplotlibで))



・28×28(epoch30、batchsize40)の潜在変数空間可視化
・256×256(epoch30、batchsize40)の潜在変数空間可視化

Comments

[ps_aux_grep]

最後の画像2枚は本当に潜在変数空間を可視化したものですか？再構成画像を可視化したようにしか見えませんが，どうなのでしょう．

また，class VAEでは train_step() で reconstruction = decoder(z) としており，VAEに渡してないdecoderを利用していると考えられます．self.decoder(z) にしてください．