Problem relevnt to PatchMerging in 3D SWIN-UNETR

[GYDDHPY]

In the 3D forwrd function of PatchMerging, it seems x3 equals to x6, which is a little bit beird.
Is this a special design for 3D image or just a code error?

MONAI/monai/networks/nets/swin_unetr.py

Lines 660 to 701 in 398466c

	class PatchMerging(nn.Module):
	"""
	Patch merging layer based on: "Liu et al.,
	Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
	<https://arxiv.org/abs/2103.14030>"
	https://github.com/microsoft/Swin-Transformer
	"""
	def __init__(self, dim: int, norm_layer: Type[LayerNorm] = nn.LayerNorm, spatial_dims: int = 3) -> None:
	"""
	Args:
	dim: number of feature channels.
	norm_layer: normalization layer.
	spatial_dims: number of spatial dims.
	"""
	super().__init__()
	self.dim = dim
	if spatial_dims == 3:
	self.reduction = nn.Linear(8 * dim, 2 * dim, bias=False)
	self.norm = norm_layer(8 * dim)
	elif spatial_dims == 2:
	self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
	self.norm = norm_layer(4 * dim)
	def forward(self, x):
	x_shape = x.size()
	if len(x_shape) == 5:
	b, d, h, w, c = x_shape
	pad_input = (h % 2 == 1) or (w % 2 == 1) or (d % 2 == 1)
	if pad_input:
	x = F.pad(x, (0, 0, 0, w % 2, 0, h % 2, 0, d % 2))
	x0 = x[:, 0::2, 0::2, 0::2, :]
	x1 = x[:, 1::2, 0::2, 0::2, :]
	x2 = x[:, 0::2, 1::2, 0::2, :]
	x3 = x[:, 0::2, 0::2, 1::2, :]
	x4 = x[:, 1::2, 0::2, 1::2, :]
	x5 = x[:, 0::2, 1::2, 0::2, :]
	x6 = x[:, 0::2, 0::2, 1::2, :]
	x7 = x[:, 1::2, 1::2, 1::2, :]
	x = torch.cat([x0, x1, x2, x3, x4, x5, x6, x7], -1)

[tangy5]

Hi Thanks for pointing out this issue. This seems to be a design option, and also a code typo here.
This issue originate whether to merge and downsample along Z axis, initially, it was

        x0 = x[:, :, 0::2, 0::2, :]  # B D H/2 W/2 C
        x1 = x[:, :, 1::2, 0::2, :]  # B D H/2 W/2 C
        x2 = x[:, :, 0::2, 1::2, :]  # B D H/2 W/2 C
        x3 = x[:, :, 1::2, 1::2, :]  # B D H/2 W/2 C
        x = torch.cat([x0, x1, x2, x3], -1)  # B D H/2 W/2 4*C

And yes, after discussion on the 3D, the patch merging should be:

            x0 = x[:, 0::2, 0::2, 0::2, :]
            x1 = x[:, 1::2, 0::2, 0::2, :]
            x2 = x[:, 0::2, 1::2, 0::2, :]
            x3 = x[:, 0::2, 0::2, 1::2, :]
            x4 = x[:, 1::2, 0::2, 1::2, :]
-           x5 = x[:, 0::2, 1::2, 0::2, :]
-           x6 = x[:, 0::2, 0::2, 1::2, :]
+           x5 = x[:, 1::2, 1::2, 0::2, :]
+           x6 = x[:, 0::2, 1::2, 1::2, :]
            x7 = x[:, 1::2, 1::2, 1::2, :]
            x = torch.cat([x0, x1, x2, x3, x4, x5, x6, x7], -1)

We added patchMergingV2 as follow:
a4ef630

Thanks again

[GYDDHPY]

Got it!
Thank you for your attention.

[function2-llx]

Maybe using a for loop is less error-prone for situation like this, I think XD

Something like:

x = torch.cat([
	x[:, i::2, j::2, k::2, :] 
	for i, j, k in itertools.product(range(2), range(2), range(2))
], -1)

[wyli]

Yes, it also improves the readability, are you interested in submitting a PR? @function2-llx

[function2-llx]

With pleasure!
It's here: #4793