Patch size cannot be changed to different size in UNETR2D (disc. #3746)

[SakaltrasNikolaos]

Hello everybody,

primarily I would like to thank you all, for the amazing work and the overall support you provide. Keep up the good work!

I am testing the UNETR2D from the discussion #3746 . I run the code as provided and it runs smoothly. But I try to test different patch sizes, but it shows an error. I do not know yet why. It happens always in the last step of concatenation. The dims do not match.

Patch sizes with the same error 32 and 64:

unetr_block.py:83, in UnetrUpBlock.forward(self, inp, skip) [80](file:///c%3A/ProgramData/Anaconda3/envs/py38/lib/site-packages/monai/networks/blocks/unetr_block.py?line=79) def forward(self, inp, skip): [81](file:///c%3A/ProgramData/Anaconda3/envs/py38/lib/site-packages/monai/networks/blocks/unetr_block.py?line=80) # number of channels for skip should equals to out_channels [82](file:///c%3A/ProgramData/Anaconda3/envs/py38/lib/site-packages/monai/networks/blocks/unetr_block.py?line=81) out = self.transp_conv(inp) ---> [83](file:///c%3A/ProgramData/Anaconda3/envs/py38/lib/site-packages/monai/networks/blocks/unetr_block.py?line=82) out = torch.cat((out, skip), dim=1) [84](file:///c%3A/ProgramData/Anaconda3/envs/py38/lib/site-packages/monai/networks/blocks/unetr_block.py?line=83) out = self.conv_block(out) [85](file:///c%3A/ProgramData/Anaconda3/envs/py38/lib/site-packages/monai/networks/blocks/unetr_block.py?line=84) return out
ERROR: RuntimeError: Sizes of tensors must match except in dimension 2. Got 256 and 64 (The offending index is 0): Patch 64x64
ERROR: RuntimeError: Sizes of tensors must match except in dimension 2. Got 256 and 128 (The offending index is 0): Patch 32x32

Is the code fixed to work only with this patch size of 16x16, or am I missing something?

I add the code again here, if there is something I am missing!

`# %%
from typing import Tuple, Union
import torch
import torch.nn as nn

from monai.networks.blocks.dynunet_block import UnetOutBlock
from monai.networks.blocks import UnetrBasicBlock, UnetrPrUpBlock, UnetrUpBlock
from monai.networks.nets import ViT

%%

class UNETR2D(nn.Module):
"""
UNETR based on: "Hatamizadeh et al.,
UNETR: Transformers for 3D Medical Image Segmentation https://arxiv.org/abs/2103.10504"
"""

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    img_size: Tuple[int, int],
    feature_size: int = 16,
    hidden_size: int = 768,
    mlp_dim: int = 3072,
    num_heads: int = 12,
    pos_embed: str = "perceptron",
    norm_name: Union[Tuple, str] = "instance",
    conv_block: bool = False,
    res_block: bool = True,
    dropout_rate: float = 0.0,
) -> None:

    super().__init__()

    if not (0 <= dropout_rate <= 1):
        raise AssertionError("dropout_rate should be between 0 and 1.")

    if hidden_size % num_heads != 0:
        raise AssertionError("hidden size should be divisible by num_heads.")

    if pos_embed not in ["conv", "perceptron"]:
        raise KeyError(f"Position embedding layer of type {pos_embed} is not supported.")

    self.num_layers = 12
    self.patch_size = (16, 16)
    # self.patch_size = (32, 32)
    self.feat_size = tuple(img_d // p_d for img_d, p_d in zip(img_size, self.patch_size))
    self.hidden_size = hidden_size
    self.classification = False
    self.vit = ViT(
        in_channels=in_channels,
        img_size=img_size,
        patch_size=self.patch_size,
        hidden_size=hidden_size,
        mlp_dim=mlp_dim,
        num_layers=self.num_layers,
        num_heads=num_heads,
        pos_embed=pos_embed,
        classification=self.classification,
        dropout_rate=dropout_rate,
        spatial_dims=2
    )
    self.encoder1 = UnetrBasicBlock(
        spatial_dims=2,
        in_channels=in_channels,
        out_channels=feature_size,
        kernel_size=3,
        stride=1,
        norm_name=norm_name,
        res_block=res_block,
    )
    self.encoder2 = UnetrPrUpBlock(
        spatial_dims=2,
        in_channels=hidden_size,
        out_channels=feature_size * 2,
        num_layer=2,
        kernel_size=3,
        stride=1,
        upsample_kernel_size=2,
        norm_name=norm_name,
        conv_block=conv_block,
        res_block=res_block,
    )
    self.encoder3 = UnetrPrUpBlock(
        spatial_dims=2,
        in_channels=hidden_size,
        out_channels=feature_size * 4,
        num_layer=1,
        kernel_size=3,
        stride=1,
        upsample_kernel_size=2,
        norm_name=norm_name,
        conv_block=conv_block,
        res_block=res_block,
    )
    self.encoder4 = UnetrPrUpBlock(
        spatial_dims=2,
        in_channels=hidden_size,
        out_channels=feature_size * 8,
        num_layer=0,
        kernel_size=3,
        stride=1,
        upsample_kernel_size=2,
        norm_name=norm_name,
        conv_block=conv_block,
        res_block=res_block,
    )
    self.decoder5 = UnetrUpBlock(
        spatial_dims=2,
        in_channels=hidden_size,
        out_channels=feature_size * 8,
        kernel_size=3,
        upsample_kernel_size=2,
        norm_name=norm_name,
        res_block=res_block,
    )
    self.decoder4 = UnetrUpBlock(
        spatial_dims=2,
        in_channels=feature_size * 8,
        out_channels=feature_size * 4,
        kernel_size=3,
        upsample_kernel_size=2,
        norm_name=norm_name,
        res_block=res_block,
    )
    self.decoder3 = UnetrUpBlock(
        spatial_dims=2,
        in_channels=feature_size * 4,
        out_channels=feature_size * 2,
        kernel_size=3,
        upsample_kernel_size=2,
        norm_name=norm_name,
        res_block=res_block,
    )
    self.decoder2 = UnetrUpBlock(
        spatial_dims=2,
        in_channels=feature_size * 2,
        out_channels=feature_size,
        kernel_size=3,
        upsample_kernel_size=2,
        norm_name=norm_name,
        res_block=res_block,
    )
    self.out = UnetOutBlock(spatial_dims=2, in_channels=feature_size, out_channels=out_channels)  # type: ignore

def proj_feat(self, x, hidden_size, feat_size):
    x = x.view(x.size(0), feat_size[0], feat_size[1], hidden_size)
    x = x.permute(0, 3, 1, 2).contiguous()
    return x

def forward(self, x_in):
    x, hidden_states_out = self.vit(x_in)
    enc1 = self.encoder1(x_in)
    x2 = hidden_states_out[3]
    enc2 = self.encoder2(self.proj_feat(x2, self.hidden_size, self.feat_size))
    x3 = hidden_states_out[6]
    enc3 = self.encoder3(self.proj_feat(x3, self.hidden_size, self.feat_size))
    x4 = hidden_states_out[9]
    enc4 = self.encoder4(self.proj_feat(x4, self.hidden_size, self.feat_size))
    dec4 = self.proj_feat(x, self.hidden_size, self.feat_size)
    dec3 = self.decoder5(dec4, enc4)
    dec2 = self.decoder4(dec3, enc3)
    dec1 = self.decoder3(dec2, enc2)
    out = self.decoder2(dec1, enc1)
    logits = self.out(out)
    return logits

%%

model = UNETR2D(
in_channels=1, # 3 channels, R,G,B
out_channels=1,
img_size=(256, 256),
feature_size=16,
hidden_size=768,
mlp_dim=3072,
num_heads=12,
pos_embed="perceptron",
norm_name="instance",
res_block=True,
dropout_rate=0.0,
).cuda()

%%

print(model)

%%

x = torch.rand((1,1,256,256)).cuda()
y = model(x)`

Thank you in advance for your time and assistance. I apologize is the discussion topic exists somewhere else, and I could not find it.

Regards,
Nikolaos

[Nic-Ma]

Hi @heyufan1995 ,

Could you please help share some best practice about this question?

Thanks in advance.

[SakaltrasNikolaos]

Thank you for your time. I appreciate it.

[heyufan1995]

Hi @SakaltrasNikolaos:
The issue is that unetr has a fixed number of decoders that upsample the image 2x each time. So it only upsamples the image by 16x. That's why the VIT patch size is 16, and matches the upsample ratio 16. If you want to change patch size, the simpliest way is to add an nn.Upsample layer in the forward pass:

def __init__(...):
        ....
        self.num_layers = 12
        img_size = ensure_tuple_rep(img_size, spatial_dims)
        self.patch_size = ensure_tuple_rep(64, spatial_dims) # change the patch size from 16 to 64
        ...
def forward(...):
        ...
        dec3 = self.decoder5(dec4, enc4)
        dec2 = self.decoder4(dec3, enc3)
        dec1 = self.decoder3(dec2, enc2)
        dec1 = nn.Upsample(scale_factor=4)(dec1) # upsample it by 4. If patch_size is 32, upsample it by 2.
        out = self.decoder2(dec1, enc1)
        ...

[SakaltrasNikolaos]

Thank you for your time and fast response!