I am trying to adapt tensorflow’s example UNet (https://www.tensorflow.org/tutorials/images/segmentation) for my purposes. The main difference is this UNet takes 128×128 images and masks and my images are 512×512 and my masks are 100×100.

I am getting this error when I try to run model.fit:

ValueError: as_list() is not defined on an unknown TensorShape.

However, I can run model.predict without issue and it generates a prediction I’d expect from an untrained model.

Here is the code I used to make and train my model:

base_model = tf.keras.applications.MobileNetV2(input_shape=[512, 512, 3], include_top=False)

# Use the activations of these layers
layer_names = [
    'block_1_expand_relu',   # 64x64
    'block_3_expand_relu',   # 32x32
    'block_6_expand_relu',   # 16x16
    'block_13_expand_relu',  # 8x8
    'block_16_project',      # 4x4
]
base_model_outputs = [base_model.get_layer(name).output for name in layer_names]

# Create the feature extraction model
down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)

down_stack.trainable = False

up_stack = [
    pix2pix.upsample(512, 3),  # 4x4 -> 8x8
    pix2pix.upsample(256, 3),  # 8x8 -> 16x16
    pix2pix.upsample(128, 3),  # 16x16 -> 32x32
    pix2pix.upsample(64, 3),   # 32x32 -> 64x64
]

def unet_model(output_channels:int):
    inputs = tf.keras.layers.Input(shape=[512, 512, 3])

    # Downsampling through the model
    skips = down_stack(inputs)
    x = skips[-1]
    skips = reversed(skips[:-1])

    # Upsampling and establishing the skip connections
    for up, skip in zip(up_stack, skips):
        x = up(x)
        concat = tf.keras.layers.Concatenate()
        x = concat([x, skip])

    # This is the last layer of the model
    last = tf.keras.layers.Conv2DTranspose(
        filters=output_channels, kernel_size=3, strides=2,
        padding='same')  #64x64 -> 128x128

    x = last(x)

    return tf.keras.Model(inputs=inputs, outputs=x)

OUTPUT_CLASSES = 5

model = unet_model(output_channels=OUTPUT_CLASSES)
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model_history = model.fit(train_dataset, epochs=EPOCHS,
                        validation_data=test_dataset)

I tried examining the image shape and mask shape of each batch. The image shape is (16, 512, 512, 3) and the mask shape is (16, 512, 512, 1) for every batch but the last (which only has one image).

I tried putting this code in my process_paths (what the tutorial called it) function:

image = tf.reshape(image, [512, 512, 3])

...
mask = tf.reshape(mask, [100, 100, 1])

I tried playing around with the numbers in up_stack part a bit but ultimately didn’t get anywhere because I don’t understand that part. My assumption is that, now that I’ve changed the input size, I have to change the output size of the model’s layers, but I’m not really sure how to do that. Also, I am confused why I can still run model.predict if that’s the case.

My tensorflow version is 2.16.1

Thanks!