多类别语义分割（UNet ++）

Question

我正在尝试使用 Unet ++ 进行语义细分。

shape: [N, H, W] i.e: [16, 256, 256]

target shape: [N, 1, H, W] i.e: [16, 1, 256, 256]

output shape: [N, 1, H, W] i.e: [16, 1, 256, 256]

loss function: BCEWithLogitsLoss

我已使用this model architecture和 deepsupervision 训练了模型。而且2堂课都很好。

现在，我正在尝试将此受训模型用于 3个班级。

我已经将我的蒙版像素值转换为相应的类值。 因此，新的目标形状变为[N, H, W] i.e: [16, 256, 256]

现在我不明白如何在 3个班级中使用相同的模型。与模型体系结构一样，out_channels = 1是固定的。

我尝试过

class_number = 3
model.final1.out_channels = class_number
model.final2.out_channels = class_number
model.final3.out_channels = class_number
model.final4.out_channels = class_number

但是似乎它仍然会生成[16, 1, 256, 256]的输出，我希望得到[16, 3, 256, 256]

我是PyTorch的新手。因此，您的建议将不胜感激。
在此先感谢您。