Слой Keras запрашивает другую форму, чем в сводке

Я пишу U-net CNN в keras и пытаюсь использовать fit_generator для обучения. Чтобы это работало, я использовал скрипт-генератор, который мог передавать изображения и метки для моей сети (функция простой подгонки работает, но я хочу обучить большой набор данных, который не помещается в память).

Моя проблема в том, что в сводке модели правильно сказано, что выходной слой имеет форму: (Нет, 288, 512, 4) https://i.imgur.com/69xG8pO.jpg

но когда я пробую настоящую тренировку, я получаю эту ошибку: https://i.imgur.com/j7H6sHX.jpg

Я не понимаю, почему keras хочет (288, 512, 1), когда в сводке он ожидает (288, 512, 4)

Я попробовал это со своим собственным кодом unet, а также скопировал рабочий код с github, но у них обоих есть одна и та же проблема, которая заставляет меня поверить, что мой сценарий генератора является слабым звеном. Ниже приведен код, который я использовал (используемые здесь функции массива изображений и меток уже работали, когда я использовал их с «подгонкой» в предыдущей CNN):

def generator(img_path, label_path, batch_size, height, width, num_classes):

    input_pairs = get_pairs(img_path, label_path) # rewrite if param name changes
    random.shuffle(input_pairs)

    iterate_pairs = itertools.cycle(input_pairs)

    while True:
        X = []
        Y = []
        for _ in range(batch_size):
            im, lab = next(iterate_pairs)

            appended_im = next(iter(im))
            appended_lab = next(iter(lab))

            X.append(input_image_array(appended_im, width, height))
            Y.append(input_label_array(appended_lab, width, height, num_classes, palette))

        yield (np.array(X), np.array(Y))

Я попробовал генератор, и предоставленные партии имеют форму (для размера партии 15): (15, 288, 512, 3) (15, 288, 512, 4) Так что я действительно не знаю, в чем может быть проблема.

Обновлено: Вот код модели, который я использовал:

def conv_block(input_tensor, n_filter, kernel=(3, 3), padding='same', initializer = "he_normal"):
    x = Conv2D(n_filter, kernel, padding=padding, kernel_initializer=initializer)(input_tensor)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)
    x = Conv2D(n_filter, kernel, padding=padding, kernel_initializer=initializer)(x)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)
    return x

def deconv_block(input_tensor, residual, n_filter, kernel=(3, 3), strides=(2, 2), padding='same'):
    y = Conv2DTranspose(n_filter, kernel, strides, padding)(input_tensor)
    y = concatenate([y, residual], axis=3)
    y = conv_block(y, n_filter)
    return y

# NETWORK - n_classes is the desired number of classes, filters are fixed
def Unet(input_height, input_width, n_classes=4, filters=64):

    # Downsampling
    input_layer = Input(shape=(input_height, input_width, 3), name='input')

    conv_1 = conv_block(input_layer, filters)
    conv_1_out = MaxPooling2D(pool_size=(2, 2))(conv_1)

    conv_2 = conv_block(conv_1_out, filters*2)
    conv_2_out = MaxPooling2D(pool_size=(2, 2))(conv_2)

    conv_3 = conv_block(conv_2_out, filters*4)
    conv_3_out = MaxPooling2D(pool_size=(2, 2))(conv_3)

    conv_4 = conv_block(conv_3_out, filters*8)
    conv_4_out = MaxPooling2D(pool_size=(2, 2))(conv_4)
    conv_4_drop = Dropout(0.5)(conv_4_out)

    conv_5 = conv_block(conv_4_drop, filters*16)
    conv_5_drop = Dropout(0.5)(conv_5)

    # Upsampling
    deconv_1 = deconv_block(conv_5_drop, conv_4, filters*8)
    deconv_1_drop = Dropout(0.5)(deconv_1)

    deconv_2 = deconv_block(deconv_1_drop, conv_3, filters*4)
    deconv_2_drop = Dropout(0.5)(deconv_2)

    deconv_3 = deconv_block(deconv_2_drop, conv_2, filters*2)
    deconv_3 = deconv_block(deconv_3, conv_1, filters)

    # Output - mapping each 64-component feature vector to number of classes
    output = Conv2D(n_classes, (1, 1))(deconv_3)
    output = BatchNormalization()(output)
    output = Activation("softmax")(output)

    # embed into functional API
    model = Model(inputs=input_layer, outputs=output, name = "Unet")
    return model