About the final prediction head

[limyunn]

Hi , when reading your code about the model head (define number of output classes) ,I have some question about this. In 2D segmentation task , we usually use Conv2D with kernel size 1 to get the final segmentaion mask. However, in this library, I notice that Conv3D with kernel size 3 has been used to get the final prediction rather than kernel size 1. I'm just wondering is there difference between Conv3D with kernel size 3 and 1 to get the prediction in 3D segmentation task?

[ZFTurbo]

Can you please send link on code with line?

[limyunn]

the code is from this script : segmentation_models_3D/models/unet.py

model head (define number of output classes)

x = layers.Conv3D(
    filters=classes,
    kernel_size=(3, 3, 3),
    padding='same',
    use_bias=True,
    kernel_initializer='glorot_uniform',
    name='final_conv',
)(x)

[ZFTurbo]

It's the same in 2D version: https://github.com/qubvel/segmentation_models/blob/master/segmentation_models/models/unet.py#L142

[limyunn]

Thanks, as I have seen, most of the prediction head in Unet have implemented Conv2D with kernel size 1. So I am not sure if the kernel size 3 or 1 will affect the final prediction. Here is the original version from the author of Unet: https://github.com/zhixuhao/unet/blob/b45af4d458437d8281cc218a07fd4380818ece4a/model.py#L53