For extracting features from images I use pretrained ResNet50 (removing the classifier head). In order to transform the model to a hierarchical model, I want to use the following strategy:
import torch
from torchvision.models import resnet50, ResNet50_Weights
import torch.nn as nn
class model(nn.Module):
def __init__(self):
super(model, self).__init__()
self.model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
self.block1 = nn.Sequential(*list(self.model.children())[:5])
self.block2 = nn.Sequential(*list(self.model.children())[5:6])
self.block3 = nn.Sequential(*list(self.model.children())[6:7])
self.block4 = nn.Sequential(*list(self.model.children())[7:8])
def forward(self, x):
x1 = self.block1(x)
x2 = self.block2(x1)
x3 = self.block3(x2)
x4 = self.block4(x3)
return(x1, x2, x3, x4)
this is because I need to have the intermediate outputs. My question is that does this approach cause any problem for assigning the pretrained weights to different layers (i.e., self.block2, self.block3 ,…)? Should I edit the source code for my goal or the mentioned strategy is a working idea?