Removing some layers in a pre-trained model in pytorch - image-processing

Question

I need to use a pytorch pre-trained model but I want to remove some layers of that. This is the original model:
def forward_features(self, x):
x = self.patch_embed1(x)
x = self.pos_drop(x)
for i, blk in enumerate(self.blocks1):
if self.use_checkpoint and i < self.checkpoint_num[0]:
x = checkpoint.checkpoint(blk, x)
else:
x = blk(x)
x = self.patch_embed2(x)
for i, blk in enumerate(self.blocks2):
if self.use_checkpoint and i < self.checkpoint_num[1]:
x = checkpoint.checkpoint(blk, x)
else:
x = blk(x)
x = self.patch_embed3(x)
for i, blk in enumerate(self.blocks3):
if self.use_checkpoint and i < self.checkpoint_num[2]:
x = checkpoint.checkpoint(blk, x)
else:
x = blk(x)
x = self.patch_embed4(x)
for i, blk in enumerate(self.blocks4):
if self.use_checkpoint and i < self.checkpoint_num[3]:
x = checkpoint.checkpoint(blk, x)
else:
x = blk(x)
x = self.norm(x)
x = self.pre_logits(x)
return x
def forward(self, x):
x = x[0] #(2)
x = self.forward_features(x)
x = x.flatten(2).mean(-1)
x = self.head(x)
return x
I want to change the model in this manner:
def forward_features(self, x):
x = self.patch_embed1(x)
x = self.pos_drop(x)
for i, blk in enumerate(self.blocks1):
if self.use_checkpoint and i < self.checkpoint_num[0]:
features1 = checkpoint.checkpoint(blk, x)
else:
features1 = blk(x)
y_2 = self.patch_embed2(features1)
for i, blk in enumerate(self.blocks2):
if self.use_checkpoint and i < self.checkpoint_num[1]:
features2 = checkpoint.checkpoint(blk, y_2)
else:
features2 = blk(y_2)
y_3 = self.patch_embed3(features2)
for i, blk in enumerate(self.blocks3):
if self.use_checkpoint and i < self.checkpoint_num[2]:
features3 = checkpoint.checkpoint(blk, y_3)
else:
features3 = blk(y_3)
y_4 = self.patch_embed4(features3)
for i, blk in enumerate(self.blocks4):
if self.use_checkpoint and i < self.checkpoint_num[3]:
features4 = checkpoint.checkpoint(blk, y_4)
else:
features4 = blk(y_4)
return x, features1, features2, features3, features4
def forward(self, x):
x, features1, features2, features3, features4 = self.forward_features(z)
return features1, features2, features3, features4
How can I make these changes without having problems of mismatch between pretrained weights and the model? This is because the pytorch does not accept my modifications when it checks the similarity between model and the wrights.
Traceback of the error:
Traceback (most recent call last):
File "/home/user/code/train.py", line 370, in <module>
main()
File "/home/user/code/train.py", line 101, in main
model = modelName(pretrained=encoder_pretrained)
File "/home/user/code/models/model.py", line 72, in __init__
assert list_model_dict[i][1].shape==list_weight_dict[i][1].shape
AssertionError
The error is faced in the following part:
if pretrained:
print('Loading weights...')
weight_dict=torch.load(os.path.join('models','uniformer_small_k400_16x4.pth'))
model_dict=self.featureExtractor.state_dict()
list_model_dict=list(model_dict.items())
list_weight_dict=list(weight_dict.items())
for i in range(len(list_model_dict)):
assert list_model_dict[i][1].shape==list_weight_dict[i][1].shape # This is the line of 72
model_dict[list_model_dict[i][0]].copy_(weight_dict[list_weight_dict[i][0]])
for i in range(len(list_model_dict)):
assert torch.all(torch.eq(model_dict[list_model_dict[i][0]],weight_dict[list_weight_dict[i][0]].to('cpu')))
print('Loading done!')