Unetr bottleneck layer weight freezing for transfer learning

[justinzhang528]

Dear MONAI Community,

I want to freeze the weights of the bottleneck layers of Unetr for transfer learning:

I want to set "requires_grad = True" for the weights of the bottleneck layers and "requires_grad = False" for other layers except the bottleneck, but i have no idea which layers are bottleneck layers and which aren't. Could anyone help me please?

Thanks in advance.

[tangy5]

Hi @justinzhang528 , you could see the forward function in UNETR. the so called bottleneck layer output should be the "x"

        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))
        x3 = hidden_states_out[6]
        enc3 = self.encoder3(self.proj_feat(x3))
        x4 = hidden_states_out[9]

I assume what you need to freeze is the self.vit, and "x" is the bottleneck feature, you can also get intermediate (hidden) features from VIT at different blocks.
Hope these helps, thanks.

[justinzhang528]

It helps me a lot, thanks.

[justinzhang528]

@tangy5 Sorry, I want to know bottleneck feature about SwinUnetr as well

Here are SwinUnetr layer weights:

No	layers
1	swinViT.patch_embed.proj.weight
2	swinViT.patch_embed.proj.bias
3	swinViT.layers1.0.blocks.0.norm1.weight
4	swinViT.layers1.0.blocks.0.norm1.bias
5	swinViT.layers1.0.blocks.0.attn.relative_position_bias_table
6	swinViT.layers1.0.blocks.0.attn.qkv.weight
7	swinViT.layers1.0.blocks.0.attn.qkv.bias
8	swinViT.layers1.0.blocks.0.attn.proj.weight
9	swinViT.layers1.0.blocks.0.attn.proj.bias
10	swinViT.layers1.0.blocks.0.norm2.weight
11	swinViT.layers1.0.blocks.0.norm2.bias
12	swinViT.layers1.0.blocks.0.mlp.linear1.weight
13	swinViT.layers1.0.blocks.0.mlp.linear1.bias
14	swinViT.layers1.0.blocks.0.mlp.linear2.weight
15	swinViT.layers1.0.blocks.0.mlp.linear2.bias
16	swinViT.layers1.0.blocks.1.norm1.weight
17	swinViT.layers1.0.blocks.1.norm1.bias
18	swinViT.layers1.0.blocks.1.attn.relative_position_bias_table
19	swinViT.layers1.0.blocks.1.attn.qkv.weight
20	swinViT.layers1.0.blocks.1.attn.qkv.bias
21	swinViT.layers1.0.blocks.1.attn.proj.weight
22	swinViT.layers1.0.blocks.1.attn.proj.bias
23	swinViT.layers1.0.blocks.1.norm2.weight
24	swinViT.layers1.0.blocks.1.norm2.bias
25	swinViT.layers1.0.blocks.1.mlp.linear1.weight
26	swinViT.layers1.0.blocks.1.mlp.linear1.bias
27	swinViT.layers1.0.blocks.1.mlp.linear2.weight
28	swinViT.layers1.0.blocks.1.mlp.linear2.bias
29	swinViT.layers1.0.downsample.reduction.weight
30	swinViT.layers1.0.downsample.norm.weight
31	swinViT.layers1.0.downsample.norm.bias
32	swinViT.layers2.0.blocks.0.norm1.weight
33	swinViT.layers2.0.blocks.0.norm1.bias
34	swinViT.layers2.0.blocks.0.attn.relative_position_bias_table
35	swinViT.layers2.0.blocks.0.attn.qkv.weight
36	swinViT.layers2.0.blocks.0.attn.qkv.bias
37	swinViT.layers2.0.blocks.0.attn.proj.weight
38	swinViT.layers2.0.blocks.0.attn.proj.bias
39	swinViT.layers2.0.blocks.0.norm2.weight
40	swinViT.layers2.0.blocks.0.norm2.bias
41	swinViT.layers2.0.blocks.0.mlp.linear1.weight
42	swinViT.layers2.0.blocks.0.mlp.linear1.bias
43	swinViT.layers2.0.blocks.0.mlp.linear2.weight
44	swinViT.layers2.0.blocks.0.mlp.linear2.bias
45	swinViT.layers2.0.blocks.1.norm1.weight
46	swinViT.layers2.0.blocks.1.norm1.bias
47	swinViT.layers2.0.blocks.1.attn.relative_position_bias_table
48	swinViT.layers2.0.blocks.1.attn.qkv.weight
49	swinViT.layers2.0.blocks.1.attn.qkv.bias
50	swinViT.layers2.0.blocks.1.attn.proj.weight
51	swinViT.layers2.0.blocks.1.attn.proj.bias
52	swinViT.layers2.0.blocks.1.norm2.weight
53	swinViT.layers2.0.blocks.1.norm2.bias
54	swinViT.layers2.0.blocks.1.mlp.linear1.weight
55	swinViT.layers2.0.blocks.1.mlp.linear1.bias
56	swinViT.layers2.0.blocks.1.mlp.linear2.weight
57	swinViT.layers2.0.blocks.1.mlp.linear2.bias
58	swinViT.layers2.0.downsample.reduction.weight
59	swinViT.layers2.0.downsample.norm.weight
60	swinViT.layers2.0.downsample.norm.bias
61	swinViT.layers3.0.blocks.0.norm1.weight
62	swinViT.layers3.0.blocks.0.norm1.bias
63	swinViT.layers3.0.blocks.0.attn.relative_position_bias_table
64	swinViT.layers3.0.blocks.0.attn.qkv.weight
65	swinViT.layers3.0.blocks.0.attn.qkv.bias
66	swinViT.layers3.0.blocks.0.attn.proj.weight
67	swinViT.layers3.0.blocks.0.attn.proj.bias
68	swinViT.layers3.0.blocks.0.norm2.weight
69	swinViT.layers3.0.blocks.0.norm2.bias
70	swinViT.layers3.0.blocks.0.mlp.linear1.weight
71	swinViT.layers3.0.blocks.0.mlp.linear1.bias
72	swinViT.layers3.0.blocks.0.mlp.linear2.weight
73	swinViT.layers3.0.blocks.0.mlp.linear2.bias
74	swinViT.layers3.0.blocks.1.norm1.weight
75	swinViT.layers3.0.blocks.1.norm1.bias
76	swinViT.layers3.0.blocks.1.attn.relative_position_bias_table
77	swinViT.layers3.0.blocks.1.attn.qkv.weight
78	swinViT.layers3.0.blocks.1.attn.qkv.bias
79	swinViT.layers3.0.blocks.1.attn.proj.weight
80	swinViT.layers3.0.blocks.1.attn.proj.bias
81	swinViT.layers3.0.blocks.1.norm2.weight
82	swinViT.layers3.0.blocks.1.norm2.bias
83	swinViT.layers3.0.blocks.1.mlp.linear1.weight
84	swinViT.layers3.0.blocks.1.mlp.linear1.bias
85	swinViT.layers3.0.blocks.1.mlp.linear2.weight
86	swinViT.layers3.0.blocks.1.mlp.linear2.bias
87	swinViT.layers3.0.downsample.reduction.weight
88	swinViT.layers3.0.downsample.norm.weight
89	swinViT.layers3.0.downsample.norm.bias
90	swinViT.layers4.0.blocks.0.norm1.weight
91	swinViT.layers4.0.blocks.0.norm1.bias
92	swinViT.layers4.0.blocks.0.attn.relative_position_bias_table
93	swinViT.layers4.0.blocks.0.attn.qkv.weight
94	swinViT.layers4.0.blocks.0.attn.qkv.bias
95	swinViT.layers4.0.blocks.0.attn.proj.weight
96	swinViT.layers4.0.blocks.0.attn.proj.bias
97	swinViT.layers4.0.blocks.0.norm2.weight
98	swinViT.layers4.0.blocks.0.norm2.bias
99	swinViT.layers4.0.blocks.0.mlp.linear1.weight
100	swinViT.layers4.0.blocks.0.mlp.linear1.bias
101	swinViT.layers4.0.blocks.0.mlp.linear2.weight
102	swinViT.layers4.0.blocks.0.mlp.linear2.bias
103	swinViT.layers4.0.blocks.1.norm1.weight
104	swinViT.layers4.0.blocks.1.norm1.bias
105	swinViT.layers4.0.blocks.1.attn.relative_position_bias_table
106	swinViT.layers4.0.blocks.1.attn.qkv.weight
107	swinViT.layers4.0.blocks.1.attn.qkv.bias
108	swinViT.layers4.0.blocks.1.attn.proj.weight
109	swinViT.layers4.0.blocks.1.attn.proj.bias
110	swinViT.layers4.0.blocks.1.norm2.weight
111	swinViT.layers4.0.blocks.1.norm2.bias
112	swinViT.layers4.0.blocks.1.mlp.linear1.weight
113	swinViT.layers4.0.blocks.1.mlp.linear1.bias
114	swinViT.layers4.0.blocks.1.mlp.linear2.weight
115	swinViT.layers4.0.blocks.1.mlp.linear2.bias
116	swinViT.layers4.0.downsample.reduction.weight
117	swinViT.layers4.0.downsample.norm.weight
118	swinViT.layers4.0.downsample.norm.bias
119	encoder1.layer.conv1.conv.weight
120	encoder1.layer.conv2.conv.weight
121	encoder1.layer.conv3.conv.weight
122	encoder2.layer.conv1.conv.weight
123	encoder2.layer.conv2.conv.weight
124	encoder3.layer.conv1.conv.weight
125	encoder3.layer.conv2.conv.weight
126	encoder4.layer.conv1.conv.weight
127	encoder4.layer.conv2.conv.weight
128	encoder10.layer.conv1.conv.weight
129	encoder10.layer.conv2.conv.weight
130	decoder5.transp_conv.conv.weight
131	decoder5.conv_block.conv1.conv.weight
132	decoder5.conv_block.conv2.conv.weight
133	decoder5.conv_block.conv3.conv.weight
134	decoder4.transp_conv.conv.weight
135	decoder4.conv_block.conv1.conv.weight
136	decoder4.conv_block.conv2.conv.weight
137	decoder4.conv_block.conv3.conv.weight
138	decoder3.transp_conv.conv.weight
139	decoder3.conv_block.conv1.conv.weight
140	decoder3.conv_block.conv2.conv.weight
141	decoder3.conv_block.conv3.conv.weight
142	decoder2.transp_conv.conv.weight
143	decoder2.conv_block.conv1.conv.weight
144	decoder2.conv_block.conv2.conv.weight
145	decoder2.conv_block.conv3.conv.weight
146	decoder1.transp_conv.conv.weight
147	decoder1.conv_block.conv1.conv.weight
148	decoder1.conv_block.conv2.conv.weight
149	decoder1.conv_block.conv3.conv.weight
150	out.conv.conv.weight
151	out.conv.conv.bias

Could you tell me a little about which layers are bottleneck? Thanks!