InvalidArgumentError при реализации вариационного автоэнкодера в кератах

Question

Я пытаюсь реализовать вариационный автоэнкодер, используя в качестве справочника книгу глубокого обучения Chollet. У меня возникают трудности с подгонкой данных к модели, форма изображения (512, 512, 3), однако это дает мне ошибку, что формы несовместимы. Мне интересно, если ошибка в лямбда-слое, реализующем вариационную потерю, поскольку изображения, которые я использую, взяты из другого набора данных, чем в книге.

InvalidArgumentError: Несовместимые фигуры: [524288] против [1572864] [[{{node custom_variational_layer_6 / logistic_loss / mul}} =

Я попытался отрегулировать форму изображения, количество каналов, потери при реконструкции и изменил среднюю функцию, которая усредняет KL и кросс-энтропийные потери, на функцию суммы только кросс-энтропийных потерь, и, похоже, ничего не работает .

def create_encoder():
    input_img = Input(shape=img_shape)
    x = layers.Conv2D(32, 3, padding='same', activation='relu')(input_img)
    x = layers.Conv2D(64, 3, padding='same', strides=(2, 2), activation='relu')(x)
    x = layers.Conv2D(64, 3, padding='same', activation='relu')(x)
    x = layers.Conv2D(32, 3, padding='same', activation='relu')(x)
    shape_before_flattening = 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)
    model = Model(input_img, [z_mean, z_log_var])
    return model, shape_before_flattening

def sample(args):
    z_mean, z_log_var = args
    eps = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim),
                         mean=0., stddev=1.)
    return z_mean + K.exp(z_log_var) * eps

encoder, shape_before_flattening = create_encoder()
input_img = encoder.inputs[0]
z_mean = encoder.outputs[0]
z_log_var = encoder.outputs[1]
z = layers.Lambda(sample)([z_mean, z_log_var])

def create_decoder():
    print(K.int_shape(z))
    decoder_input = Input(K.int_shape(z)[1:])
    x = layers.Dense(np.prod(shape_before_flattening[1:]),
                     activation='relu')(decoder_input)
    x = layers.Reshape(shape_before_flattening[1:])(x)
    x = layers.Conv2DTranspose(32, 3,
                               padding='same',
                               activation='relu',
                               strides=(2, 2))(x)
    x = layers.Conv2D(1, 3,
                      padding='same',
                      activation='sigmoid')(x)
    model = Model(decoder_input, x)
    return model

class CustomVariationalLayer(keras.layers.Layer):
    def vae_loss(self, x, z_decoded): 
        x = K.flatten(x)
        z_decoded = K.flatten(z_decoded)
        xent_loss = keras.metrics.binary_crossentropy(x, z_decoded)
        kl_loss = -5e-4 * K.mean(
            1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1) 
        print(K.int_shape(xent_loss))
        print(K.int_shape(kl_loss))
        #return K.mean(mse_loss + kl_loss)
        return K.mean(xent_loss + kl_loss)

    def call(self, inputs):
        x = inputs[0]
        z_decoded = inputs[1]
        loss = self.vae_loss(x, z_decoded)
        self.add_loss(loss, inputs=inputs)
        return x

decoder = create_decoder()
z_decoded = decoder(z)
out = CustomVariationalLayer()([input_img, z_decoded])

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

X_train = X_train.astype('float32')
X_test = X_test.astype('float32')
print("X train shape ", X_train.shape)
print("X test shape ", X_test.shape)

vae = Model(input_img, out)
vae.compile(optimizer='rmsprop', loss=None)
vae.summary()
vae.fit(x=X_train, y=None, shuffle=True, epochs=10, batch_size=batch_size)

Я ожидаю, что она обучит модель и не выдаст ошибку.