Hi, I tried to return the output from an intermediate layer of a pretrained model. Have searched many times but most of the examples are for model that’s built with Sequential() method or from model Zoo. My model is a customized model (no Sequential is used here) and when I called model._modules, it only returns an OrderDict, but the order of operation is not the same as during forward(). So directly calling module(x) doesn’t work! My question is how to do this for model that’s not neither Sequential() nor VGG (models that are already in model zoo)?
Here’s my code but it doesn’t work:
class SelectiveLayer(nn.Module):
def __init__(self, to_select, model):
super(SelectiveLayer, self).__init__()
for key, module in model._modules.items():
self.add_module(key, module)
self._to_select = to_select
def forward(self, x):
list = []
for name, module in self._modules.items():
print(name, module, x.shape)
x = module(x)
if name in self._to_select:
list.append(x)
return list
class PerceptualLoss(nn.Module):
def __init__(self, pretrained_model, params):
super(PerceptualLoss, self).__init__()
self.pretrained_model = pretrained_model
self.features = SelectiveLayer("conv5", pretrained_model)
def forward(self, x):
x = self.features(x)
return x