Missing encode step in swin unetr model source code?

[isaiahchua]

Hi all,

I went through the Swin UNETR model source code and found in the forward pass of the SwinUNETR class the hidden states from 3rd stage did not pass through a UnetrBasicBlock like the other hidden states (see line 307 for swin_unetr.py). Am I missing something, is this how it's supposed to be implemented? Thanks!

def forward(self, x_in):
hidden_states_out = self.swinViT(x_in, self.normalize)
enc0 = self.encoder1(x_in)
enc1 = self.encoder2(hidden_states_out[0])
enc2 = self.encoder3(hidden_states_out[1])
enc3 = self.encoder4(hidden_states_out[2])
dec4 = self.encoder10(hidden_states_out[4])
dec3 = self.decoder5(dec4, hidden_states_out[3]) # this line
dec2 = self.decoder4(dec3, enc3)
dec1 = self.decoder3(dec2, enc2)
dec0 = self.decoder2(dec1, enc1)
out = self.decoder1(dec0, enc0)
logits = self.out(out)
return logits

[KumoLiu]

Hi @tangy5, could you please share some comments on this one?
Thanks in advance!

[EmreTaha]

Yes, we realized the same, it is very weird, the produced network is not same as the paper...

[tangy5]

Hi @EmreTaha , thanks for your interest of the work. Yes, this is intentionally implemented to keep 4 times of downsample and upsample. This is an architecture design option, if you connect 3rd stage to a UnetResblock, it will also be good but will result in a higher number of parameters with 5 times downsample and upsample. I don't think it will impact performance too much. Thanks.