Tensorflow InvalidArgumentError: Вход для изменения формы является тензором со значениями 2560, но запрошенная форма имеет 128 [Op: Reshape]

Question

Я играю с WGAN, следуя этому коду , где был использован mnist (28x28)

Я хочу изменить его для работы с изображениями в градациях серого (которые я сохранил как 1d в таблице) однако я получаю эту ошибку из сети дискриминатора, и я не совсем понимаю, какой слой его выбрасывает, так как число 2560 изменяется в зависимости от размера пакета (??).

При чтении this кажется, что ошибка зависит от того факта, что я передаю список вместо нескольких. Поэтому, даже если я вручную создаю набор данных, он ведет себя как список?

Это функции предварительной обработки, которые я использую

# datasets creation
def preprocess_():
    import pandas as pd
    osw = pd.read_csv('features.txt', sep="\t").values
    ids = osw[:, 0]
    # add 1 for greyscale
    frag = osw[:, 1:].reshape(osw.shape[0], 32, 64, 1)
    fr = tf.data.Dataset.from_tensor_slices(frag.astype(np.float32))
    lb = tf.data.Dataset.from_tensor_slices(np.ones(ids.shape))
    return fr, lb

def normalize_dataset(ex, mode='uniform'):
    ex = ex * 2 - 1.0
    ex = tf.cast(ex, tf.float32)
    #g_output, (batch_size, *image_size)
    if mode == 'uniform':
        input_z = tf.random.uniform(
            shape=(z_size, ), minval=-1.0, maxval=1.0)
    elif mode == 'normal':
        input_z = tf.random.normal(shape=(z_size, ))
    return input_z, ex

frag, y = preprocess_()
X = frag.map(normalize_dataset)

И если я проверяю размеры

image_size = (32, 64)
num_epochs = 100
batch_size = 128
z_size = 20
mode_z = 'uniform'
lambda_gp = 10.0
X = X.batch(batch_size, drop_remainder=True)
input_z, input_real = next(iter(X))
print('latent', input_z.shape)
print('real', input_real.shape)
input-z: (128, 20)
input-real -- shape: (128, 32, 64, 1)

Для wgan я изменил размер ввода, чтобы он соответствовал оттенкам серого 32 x 64

def make_dcgan_generator(
        z_size=20,
        output_size=(32, 64, 1),
        n_filters=128,
        n_blocks=2):
    size_factor = 2**n_blocks
    hidden_size = (
        output_size[0]//size_factor,
        output_size[1]//size_factor
    )

    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(z_size,)),

        tf.keras.layers.Dense(
            units=n_filters*np.prod(hidden_size),
            use_bias=False),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.LeakyReLU(),
        tf.keras.layers.Reshape(
            (hidden_size[0], hidden_size[1], n_filters)),
        tf.keras.layers.Conv2DTranspose(
            filters=n_filters, kernel_size=(5, 5), strides=(1, 1),
            padding='same', use_bias=False),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.LeakyReLU()
    ])

    nf = n_filters
    for i in range(n_blocks):
        nf = nf // 2
        model.add(
            tf.keras.layers.Conv2DTranspose(
                filters=nf, kernel_size=(5, 5), strides=(2, 2),
                padding='same', use_bias=False))
        model.add(tf.keras.layers.BatchNormalization())
        model.add(tf.keras.layers.LeakyReLU())

    model.add(
        tf.keras.layers.Conv2DTranspose(
            filters=output_size[2], kernel_size=(5, 5),
            strides=(1, 1), padding='same', use_bias=False,
            activation='tanh'))

    return model

def make_dcgan_discriminator(
        input_size=(32, 64, 1),
        n_filters=64,
        n_blocks=2):
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=input_size),
        tf.keras.layers.Conv2D(
            filters=n_filters, kernel_size=5,
            strides=(1, 1), padding='same'),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.LeakyReLU()
    ])

    nf = n_filters
    for i in range(n_blocks):
        nf = nf*2
        model.add(
            tf.keras.layers.Conv2D(
                filters=nf, kernel_size=(5, 5),
                strides=(2, 2),padding='same'))
        model.add(tf.keras.layers.BatchNormalization())
        model.add(tf.keras.layers.LeakyReLU())
        model.add(tf.keras.layers.Dropout(0.3))

    model.add(tf.keras.layers.Conv2D(
            filters=1, kernel_size=(7, 7), padding='valid'))

    model.add(tf.keras.layers.Reshape((1,)))

    return model

А обучение l oop (до исключения) следующее

with tf.device(device_name):
    gen_model = make_dcgan_generator()
    gen_model.build(input_shape=(None, z_size))

    disc_model = make_dcgan_discriminator()
    disc_model.build(input_shape=(None, np.prod(image_size)))
    disc_model.summary()


g_optimizer = tf.keras.optimizers.Adam(0.0002)
d_optimizer = tf.keras.optimizers.Adam(0.0002)
all_losses = []
epoch_samples = []

start_time = time.time()

for epoch in range(1, num_epochs+1):
    epoch_losses = []
    for i,(input_z,input_real) in enumerate(X):

        ## Compute discriminator's loss and gradients:
        with tf.GradientTape() as d_tape, tf.GradientTape() as g_tape:
            # gen_model works
            g_output = gen_model(input_z, training=True)
            d_critics_real = disc_model(input_real, training=True)
            d_critics_fake = disc_model(g_output, training=True)

Так откуда берется 2560? 32 * 64 - это 2048, 128 * 32 или 64> 2560, а внутри сети. Похоже, что нигде нет. Кроме того, при изменении размера пакета на 64 это число меняется, поэтому кажется, что оно связано с первым измерением, но не может обернуть его вокруг него.