问keras拟合函数中的尺寸错误(keras中的条件变分自动编码器)

EN

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

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

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon


class cVAE(keras.Model):
    def __init__(self, 
                 original_dim,
                 label_dim,
                 latent_dim,
                 beta=1,
                 batch_size=1,
                 **kwargs):
        super(cVAE, self).__init__(**kwargs)

        self.original_dim = original_dim
        self.latent_dim = latent_dim
        self.label_dim = label_dim
        self.beta = beta
        self.batch_size = batch_size
        # Build the encoder
        print("building encoder")
        rnaseq_inputs = keras.Input(shape=(self.original_dim, ),batch_size=self.batch_size)
        label_inputs = keras.Input(shape=(self.label_dim, ),batch_size=self.batch_size)

        encoder_inputs = layers.concatenate([rnaseq_inputs, label_inputs], name='concat_1')
        z_mean = layers.Dense(self.latent_dim, 
                              kernel_initializer = 'glorot_uniform')(encoder_inputs)
        z_mean = layers.BatchNormalization()(z_mean)
        z_mean = layers.Activation('relu')(z_mean)

        z_log_var = layers.Dense(self.latent_dim, 
                              kernel_initializer = 'glorot_uniform')(encoder_inputs)
        z_log_var = layers.BatchNormalization()(z_log_var)
        z_log_var = layers.Activation('relu')(z_log_var)
        z = Sampling()([z_mean, z_log_var])
        zc = layers.concatenate([z, label_inputs],name='concat_2')
        self.encoder = keras.Model([rnaseq_inputs, label_inputs], [z_mean, z_log_var, z, zc])
        
        print("building decoder")
        # Build the decoder
        decoder_input_dim = self.latent_dim + self.label_dim
        decoder_output_dim = self.original_dim + self.label_dim
        decoder_inputs = keras.Input(shape=(decoder_input_dim, ))
        decoder_outputs = keras.layers.Dense(decoder_output_dim, 
                                             activation='sigmoid')(decoder_inputs)
        self.decoder = keras.Model(decoder_inputs, decoder_outputs)

        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:
            # exp_data, label_data = data
            z_mean, z_log_var, z, zc = self.encoder(data)

            reconstruction = self.decoder(zc)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(
                    keras.losses.mean_squared_error(data, reconstruction)
                )
            )
            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
            total_loss = reconstruction_loss +  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 {
                "loss": self.total_loss_tracker.result(),
                "reconstruction_loss": self.reconstruction_loss_tracker.result(),
                "kl_loss": self.kl_loss_tracker.result(),
            }

toy_data = np.random.random((100,100)).astype('float32')
label = np.random.randint(0,high=2,size=100).reshape(100,1).astype('float32')

cvae_model = cVAE(original_dim=100,batch_size=2,label_dim=1,latent_dim=1)
cvae_model.compile(optimizer=keras.optimizers.Adam(learning_rate=0.003))

# fitting
cvae_model.fit([toy_data,label])

ValueError: in user code:

    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:805 train_function  *
        return step_function(self, iterator)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:795 step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/distribute_lib.py:1259 run
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/distribute_lib.py:2730 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/distribute_lib.py:3417 _call_for_each_replica
        return fn(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:788 run_step  **
        outputs = model.train_step(data)
    <ipython-input-232-1cc639e2055c>:182 train_step
        keras.losses.mean_squared_error(data, reconstruction)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py:201 wrapper
        return target(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/losses.py:1197 mean_squared_error
        y_true = math_ops.cast(y_true, y_pred.dtype)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py:201 wrapper
        return target(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/math_ops.py:964 cast
        x = ops.convert_to_tensor(x, name="x")
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/profiler/trace.py:163 wrapped
        return func(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py:1540 convert_to_tensor
        ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/array_ops.py:1525 _autopacking_conversion_function
        return _autopacking_helper(v, dtype, name or "packed")
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/array_ops.py:1444 _autopacking_helper
        converted_elem = _autopacking_helper(elem, dtype, str(i))
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/array_ops.py:1461 _autopacking_helper
        return gen_array_ops.pack(elems_as_tensors, name=scope)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/gen_array_ops.py:6398 pack
        "Pack", values=values, axis=axis, name=name)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/op_def_library.py:750 _apply_op_helper
        attrs=attr_protos, op_def=op_def)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py:592 _create_op_internal
        compute_device)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py:3536 _create_op_internal
        op_def=op_def)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py:2016 __init__
        control_input_ops, op_def)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py:1856 _create_c_op
        raise ValueError(str(e))

    ValueError: Dimension 1 in both shapes must be equal, but are 100 and 1. Shapes are [2,100] and [2,1].
        From merging shape 0 with other shapes. for '{{node Cast/x/0}} = Pack[N=2, T=DT_FLOAT, axis=0](IteratorGetNext, IteratorGetNext:1)' with input shapes: [2,100], [2,1].