使用Keras DeepLabV3+训练时sparse_categorical_crossentropy报错求助

背景 #

基于Keras的DeepLabV3+示例，针对机器人手术图像做多类语义分割，目标是分割手术器械的杆部、关节、尖端及背景。

修改的核心参数 #

• IMAGE_SIZE = 1024
• NUM_CLASSES = 4（杆部、关节、尖端、背景）
• NUM_TRAIN = 216
• NUM_VAL = 54

训练代码 #

loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.001),
    loss=loss,
    metrics=["accuracy"]
)

history = model.fit(train_dataset, validation_data=val_dataset, epochs=25)

报错信息 #

训练启动后输出：

Epoch 1/25
"--- 大型整数数组 ----"
[[{{node sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits}}]] [Op:__inference_train_function_165277]

收到错误：

Invalid Argument Error: Graph Execution error: Detected at node sparse_categorical_crossentropy/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits

已尝试的修改 #

• 将损失函数的from_logits改为False
• 将输入图像尺寸改回示例中的512

本人是Python和计算机视觉新手，求各位帮忙解决，谢谢！

内容的提问来源于stack exchange，提问作者Hardik Patel