keras 写一个高清的输入是91 109 91的VAE

以下是一个基于Keras的高清图像VAE的示例代码，输入大小为91 x 109 x 91。

import numpy as np
from keras.layers import Input, Dense, Lambda, Flatten, Reshape
from keras.layers.convolutional import Conv3D, Conv3DTranspose
from keras.models import Model
from keras import backend as K
from keras.datasets import mnist
from keras.losses import mse, binary_crossentropy

# dimensions of the input image
img_rows, img_cols, img_depth = 91, 109, 91
input_shape = (img_rows, img_cols, img_depth, 1)
latent_dim = 2

# VAE model = encoder + decoder
# build encoder model
inputs = Input(shape=input_shape, name='encoder_input')
x = Conv3D(32, kernel_size=(3, 3, 3), activation='relu', padding='same')(inputs)
x = Conv3D(64, kernel_size=(3, 3, 3), activation='relu', padding='same', strides=(2, 2, 2))(x)
x = Conv3D(64, kernel_size=(3, 3, 3), activation='relu', padding='same')(x)
x = Conv3D(64, kernel_size=(3, 3, 3), activation='relu', padding='same')(x)
shape_before_flattening = K.int_shape(x)
x = Flatten()(x)
x = Dense(32, activation='relu')(x)

z_mean = Dense(latent_dim, name='z_mean')(x)
z_log_var = Dense(latent_dim, name='z_log_var')(x)

# reparameterization trick
def sampling(args):
    z_mean, z_log_var = args
    batch = K.shape(z_mean)[0]
    dim = K.int_shape(z_mean)[1]
    epsilon = K.random_normal(shape=(batch, dim))
    return z_mean + K.exp(0.5 * z_log_var) * epsilon

# use reparameterization trick to push the sampling out as input
# note that "output_shape" isn't necessary with the TensorFlow backend
z = Lambda(sampling, output_shape=(latent_dim,), name='z')([z_mean, z_log_var])

# instantiate encoder model
encoder = Model(inputs, [z_mean, z_log_var, z], name='encoder')
encoder.summary()

# build decoder model
latent_inputs = Input(shape=(latent_dim,), name='z_sampling')
x = Dense(32, activation='relu')(latent_inputs)
x = Dense(np.prod(shape_before_flattening[1:]), activation='relu')(x)
x = Reshape(shape_before_flattening[1:])(x)
x = Conv3DTranspose(64, kernel_size=(3, 3, 3), activation='relu', padding='same')(x)
x = Conv3DTranspose(64, kernel_size=(3, 3, 3), activation='relu', padding='same')(x)
x = Conv3DTranspose(32, kernel_size=(3, 3, 3), activation='relu', padding='same', strides=(2, 2, 2))(x)
outputs = Conv3DTranspose(1, kernel_size=(3, 3, 3), activation='sigmoid', padding='same')(x)

# instantiate decoder model
decoder = Model(latent_inputs, outputs, name='decoder')
decoder.summary()

# instantiate VAE model
outputs = decoder(encoder(inputs)[2])
vae = Model(inputs, outputs, name='vae')

# VAE loss = mse_loss or xent_loss + kl_loss
reconstruction_loss = mse(K.flatten(inputs), K.flatten(outputs))
reconstruction_loss *= img_rows * img_cols * img_depth
kl_loss = 1 + z_log_var - K.square(z_mean) - K.exp(z_log_var)
kl_loss = K.sum(kl_loss, axis=-1)
kl_loss *= -0.5
vae_loss = K.mean(reconstruction_loss + kl_loss)
vae.add_loss(vae_loss)
vae.compile(optimizer='adam')
vae.summary()

# load MNIST dataset
(x_train, _), (x_test, _) = mnist.load_data()

# preprocess the data
x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.
x_train = np.reshape(x_train, (len(x_train), img_rows, img_cols, img_depth, 1))
x_test = np.reshape(x_test, (len(x_test), img_rows, img_cols, img_depth, 1))

# train the VAE model
vae.fit(x_train, epochs=10, batch_size=128, validation_data=(x_test, None))