RuntimeError при использовании keras 3 с бэкэндом pytorch ⇐ Python

[Anonymous]

Я пытаюсь воспроизвести модель VAE, продемонстрированную в документе Keras3:
https://keras.io/examples/generative/vae/ и https://keras.io/guides/making_new_laye ... bclassing/# пример-все-все-в-конец
Второй пример работает хорошо. Поскольку я пытаюсь использовать серверную часть torch, я изменил некоторые места в первом примере:

Код: Выделить всё

import os
os.environ['KERAS_BACKEND'] = 'torch'
import keras

import numpy as np
from keras import layers, ops

class Sampling(layers.Layer):
"""Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""

def __init__(self, name='sampling', **kwargs):
super(Sampling, self).__init__(name=name, **kwargs)
self.seed_generator = keras.random.SeedGenerator(42)

def call(self, inputs):
z_mean, z_log_var = inputs
batch = ops.shape(z_mean)[0]
dim = ops.shape(z_mean)[1]
epsilon = keras.random.normal(shape=(batch, dim), seed=self.seed_generator)
return z_mean + ops.exp(0.5 * z_log_var) * epsilon

class Encoder(keras.Model):
"""Maps MNIST digits to a triplet (z_mean, z_log_var, z)."""

def __init__(self, latent_dim=32, intermediate_dim=64, name='encoder', **kwargs):
super().__init__(name=name, **kwargs)
self.conv_layer1 = layers.Conv2D(32, 3, activation='relu', strides=2, padding='same')
self.conv_layer2 = layers.Conv2D(64, 3, activation='relu', strides=2, padding='same')
self.flatten = layers.Flatten()
self.dense_proj = layers.Dense(intermediate_dim, activation='relu')
self.dense_mean = layers.Dense(latent_dim, name='z_mean')
self.dense_log_var = layers.Dense(latent_dim, name='z_log_var')
self.sampling = Sampling()

def call(self, inputs):
# encoder_inputs = keras.Input(shape=(28, 28, 1))(inputs)
x = self.conv_layer1(inputs)
x = self.conv_layer2(x)
x = self.flatten(x)
x = self.dense_proj(x)
z_mean = self.dense_mean(x)
z_log_var = self.dense_log_var(x)
z = self.sampling((z_mean, z_log_var))
return z_mean, z_log_var, z

class Decoder(keras.Model):
"""Converts z, the encoded digit vector, back into a readable digit."""

def __init__(self, original_dim, intermediate_dim=64, name='decoder', **kwargs):
super(Decoder, self).__init__(name=name, **kwargs)
self.dense_proj = layers.Dense(7 * 7 * 64, activation='relu')
self.reshape = layers.Reshape((7, 7, 64))
self.conv_transpose1 = layers.Conv2DTranspose(64, 3, activation='relu', strides=2, padding='same')
self.conv_transpose2 = layers.Conv2DTranspose(32, 3, activation='relu', strides=2, padding='same')
self.dense_output = layers.Conv2DTranspose(1, 3, activation='sigmoid', padding='same')
# self.dense_output = layers.Dense(original_dim, activation='sigmoid')

def call(self, inputs):
x = self.dense_proj(inputs)
x = self.reshape(x)
x = self.conv_transpose1(x)
x = self.conv_transpose2(x)
return self.dense_output(x)

class VAE(keras.Model):
"""Combines the encoder and decoder into an end-to-end model for training."""

def __init__(
self,
encoder,
decoder,
name='vae',
**kwargs
):
super().__init__(name=name, **kwargs)
self.encoder = encoder
self.decoder = decoder

def call(self, input):
z_mean, z_log_var, z = self.encoder(input)
reconstructed = self.decoder(z)
# Add KL divergence regularization loss.
kl_loss = -0.5 * ops.mean(1 + z_log_var - ops.square(z_mean) - ops.exp(z_log_var), axis=1)
self.add_loss(kl_loss)
return reconstructed

if __name__ == '__main__':
(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()
print(x_train.shape)
mnist_digits = np.concatenate([x_train, x_test], axis=0)
mnist_digits = np.expand_dims(mnist_digits, -1).astype("float32") / 255
# x_train = x_train.reshape(-1, 28, 28, 1).astype('float32') / 255
print(mnist_digits.shape)

encoder = Encoder(latent_dim=32)
# print(encoder.summary())
decoder = Decoder(original_dim=784)
# print(decoder.summary())
vae = VAE(encoder=encoder, decoder=decoder)
vae.compile(optimizer=keras.optimizers.Adam(learning_rate=1e-3), loss=keras.losses.MeanSquaredError())
vae.fit(mnist_digits, mnist_digits, epochs=30,  batch_size=128)

Однако я постоянно получаю следующую ошибку:

Код: Выделить всё

/opt/miniconda3/envs/vae/bin/python /Users/belter/github/VAE/example.py
(60000, 28, 28)
(70000, 28, 28, 1)
/opt/miniconda3/envs/vae/lib/python3.11/site-packages/keras/src/backend/common/backend_utils.py:89: UserWarning: You might experience inconsistencies across backends when calling conv transpose with kernel_size=3, stride=2, dilation_rate=1, padding=same, output_padding=1.
warnings.warn(
Epoch 1/30
Traceback (most recent call last):
File "/Users/belter/github/VAE/example.py", line 106, in 
vae.fit(mnist_digits, mnist_digits, epochs=30, batch_size=128)
File "/opt/miniconda3/envs/vae/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/opt/miniconda3/envs/vae/lib/python3.11/site-packages/keras/src/backend/torch/numpy.py", line 1248, in stack
return torch.stack(x, dim=axis)
^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: stack expects each tensor to be equal size, but got [] at entry 0 and [128] at entry 1

Process finished with exit code 1

Я использую Python 3.11, pytorch 2.2.2 и keras 3.2.1.
Я знаю, что это может быть связано с проблемой размерность входных данных, но я не могу в этом разобраться.

Подробнее здесь: https://stackoverflow.com/questions/783 ... ch-backend