UNet diffrent stage output

[YingdiZhang]

Thanks for the great job. As I know, there are 4 downsample stage in UNet. How to output the four stage feature separately.

[KumoLiu]

Hi @YingdiZhang, maybe you could refer to this discussion.

Hope it can help you, thanks!

[YingdiZhang]

Thanks for your reply, it seems I should input an argument which is the name of intermediate layer for example "last_layer". While after reading the source code of UNet, I dont find the four stage intermediate layer's name.

[KumoLiu]

Hi @YingdiZhang, "last_layer" is just a key used in "intermediate_outputs". You could see the whole model architecture by using print(net). You could also get the intermediate output like print(net.model[1].submodule[2]).

Hope it can help you, thanks!