Can't get multi-output CNN to work (tensorflow and keras) - tensorflow

Question

I'm currently working on a task of fiber tip tracking on an endoscopic video.
For this purpose I have two models:
classifier that tells whether image contains fiber (is_visible)
regressor that predicts fiber tip position (x, y)
I am using ResNet18 pretrained on ImageNet for this purpose and it works great. But I'm experiencing performance issues,
so I decided to combine these two models into a single one using multi-output approach.
But so far I haven't been able to get it to work.
TENSORFLOW:
TensorFlow version: 2.10.1
DATATSET:
My dataset is stored in a HDF5 format. Each sample has:
an image (224, 224, 3)
uint8 for visibility flag
and two floats for fiber tip position (x, y)
I am loading this dataset using custom generator as follows:
output_types = (tf.float32, tf.uint8, tf.float32)
output_shapes = (
tf.TensorShape((None, image_height, image_width, number_of_channels)), # image
tf.TensorShape((None, 1)), # is_visible
tf.TensorShape((None, 1, 1, 2)), # x, y
)
train_dataset = tf.data.Dataset.from_generator(
generator, output_types=output_types, output_shapes=output_shapes,
)
MODEL:
My model is defined as follows:
model = ResNet18(input_shape=(224, 224, 3), weights="imagenet", include_top=False)
inputLayer = model.input
innerLayer = tf.keras.layers.Flatten()(model.output)
is_visible = tf.keras.layers.Dense(1, activation="sigmoid", name="is_visible")(innerLayer)
position = tf.keras.layers.Dense(2)(innerLayer)
position = tf.keras.layers.Reshape((1, 1, 2), name="position")(position)
model = tf.keras.Model(inputs=[inputLayer], outputs=[is_visible, position])
adam = tf.keras.optimizers.Adam(1e-4)
model.compile(
optimizer=adam,
loss={
"is_visible": "binary_crossentropy",
"position": "mean_squared_error",
},
loss_weights={
"is_visible": 1.0,
"position": 1.0
},
metrics={
"is_visible": "accuracy",
"position": "mean_squared_error"
},
)
PROBLEM:
Dataset is working great, I can loop through each batch. But when it comes to training
model.fit(
train_dataset,
validation_data=validation_dataset,
epochs=100000,
callbacks=callbacks,
)
I get the following error
ValueError: Can not squeeze dim[3], expected a dimension of 1, got 2 for '{{node mean_squared_error/weighted_loss/Squeeze}} = SqueezeT=DT_FLOAT, squeeze_dims=[-1]' with input shapes: [?,1,1,2].
I tried to change the dataset format like so:
output_types = (tf.float32, tf.uint8, tf.float32, tf.float32)
output_shapes = (
tf.TensorShape((None, image_height, image_width, number_of_channels)), # image
tf.TensorShape((None, 1)), # is_visible
tf.TensorShape((None, 1)), # x
tf.TensorShape((None, 1)), # y
)
But these leads to another error:
ValueError: Data is expected to be in format x, (x,), (x, y), or (x, y, sample_weight), found: (<tf.Tensor 'IteratorGetNext:0' shape=(None, 224, 224, 3) dtype=float32>, <tf.Tensor 'IteratorGetNext:1' shape=(None, 1) dtype=uint8>, <tf.Tensor 'IteratorGetNext:2' shape=(None, 1) dtype=float32>, <tf.Tensor 'IteratorGetNext:3' shape=(None, 1) dtype=float32>)
I tried to wrap is_visible and (x,y) returned from train_dataset into dictionary like so:
yield image_batch, {"is_visible": is_visible_batch, "position": position_batch}
Also tried these options:
yield image_batch, (is_visible_batch, position_batch)
yield image_batch, [is_visible_batch, position_batch]
But that didn't help
Can anyone tell me what am I doing wrong? I am totally stuck ))