Segmentation of one dataset with 4 different CNN models (UNet, DynUNet, SegResNet, and VNet)

[Kiarashdnsh]

Hello,

I'm a novice in AI and image segmentation.
I have a dataset of [256, 256, 256]-sized images and labels. The input channel is 1 and the output channel is 3 (label IDs are 0, 1, and 2, respectively). I'm trying to segment the images using UNet, DynUNet, VNet, and SegResNet with the MONAI framework.
For UNet, I have used the following code:

model = UNet(
        spatial_dims=3,
        in_channels=1,
        out_channels=3,
        channels=(16, 32, 64, 128, 256),
        strides=(2, 2, 2, 2),
        num_res_units=2,
        norm=Norm.BATCH,
    ).to(device)
loss_function = DiceLoss(include_background=False, to_onehot_y=True, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)

and used this tutorial here for running the trainer.

Thus far, the results have been satisfactory on the UNet front! However, I also wanted to use the same dataset with the same dataloader, but with a different network model in this framework, starting with DynUNet. To achieve this, I changed the model to:

model = DynUNet(
        spatial_dims=3,
        in_channels=1,
        out_channels=3,
        kernel_size=(3, 3, 3),
        strides=(2, 2, 2),
        upsample_kernel_size=(2, 2, 2),
        deep_supervision=False,
        filters=(2, 2, 2),
        dropout=0.2,
        norm_name="instance",
        act_name="leakyrelu",
        deep_supr_num=1,
        res_block=True,
        trans_bias=True,
)
lloss_function = DiceLoss(include_background=False, to_onehot_y=True, sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)

But now I get this error while running the same trainer code:

AssertionError: ground truth has different shape (torch.Size([2, 2, 256, 256, 256])) from input (torch.Size([2, 2, 128, 128, 128]))

• Can you help me solve this issue?
• What can I do to issuelessly implement other CNN models in the same environment?

Thanks in advance,
Kiarash

[KumoLiu]

Hi @Kiarashdnsh, thanks for your interest here.

• In your case, you downsample the image twice and upsample only once, so your output shape is different from the label. You can simply update the stride to (1, 2, 2).
  And for more details about how to set the parameters of the DynUNet, you can refer to
  

  MONAI/monai/networks/nets/dynunet.py

  Lines 70 to 85 in 17c1e3a

  	The model supports 2D or 3D inputs and is consisted with four kinds of blocks:
  	one input block, `n` downsample blocks, one bottleneck and `n+1` upsample blocks. Where, `n>0`.
  	The first and last kernel and stride values of the input sequences are used for input block and
  	bottleneck respectively, and the rest value(s) are used for downsample and upsample blocks.
  	Therefore, pleasure ensure that the length of input sequences (``kernel_size`` and ``strides``)
  	is no less than 3 in order to have at least one downsample and upsample blocks.
  	To meet the requirements of the structure, the input size for each spatial dimension should be divisible
  	by the product of all strides in the corresponding dimension. In addition, the minimal spatial size should have
  	at least one dimension that has twice the size of the product of all strides.
  	For example, if `strides=((1, 2, 4), 2, 2, 1)`, the spatial size should be divisible by `(4, 8, 16)`,
  	and the minimal spatial size is `(8, 8, 16)` or `(4, 16, 16)` or `(4, 8, 32)`.
  	The output size for each spatial dimension equals to the input size of the corresponding dimension divided by the
  	stride in strides[0].
  	For example, if `strides=((1, 2, 4), 2, 2, 1)` and the input size is `(64, 32, 32)`, the output size is `(64, 16, 8)`.

• You can replace the networks in this pipeline with any network seamlessly, but you need to at least understand the structure of each network and know how to set the parameters to get the desired output size.

Hope it can help you, thanks!

[Kiarashdnsh]

Thanks. With your help, I managed to solve the issue. But where do I find more info on the structure of different models and the functionality of their hyper/parameters? The Monai's documentation is not detailed enough for someone at my level. For example, I don't know the what the sw_batch_size is in the training pipeline, and how it is connected with the roi_size parameter, and the documentation doesn't provide in-depth info on it.
I need information on many of these parameters and their functionality. Thanks! 🙌

[KumoLiu]

Hi @Kiarashdnsh, you can print the network architecture or maybe just read the original paper we cited. I think most of the models in MONAI have citations like these.

MONAI/monai/networks/nets/dynunet.py

Lines 58 to 61 in 17c1e3a

	This reimplementation of a dynamic UNet (DynUNet) is based on:
	`Automated Design of Deep Learning Methods for Biomedical Image Segmentation <https://arxiv.org/abs/1904.08128>`_.
	`nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation <https://arxiv.org/abs/1809.10486>`_.
	`Optimized U-Net for Brain Tumor Segmentation <https://arxiv.org/pdf/2110.03352.pdf>`_.

Thanks!