In order to extract intermediate features from a torchvision model, I implemented a class where the forward return the ouput of the network and the intermediate features.
This is what it looks like:
import torch
import torchvision.models as models
import torch.nn as nn
from copy import deepcopy
class CustomMobilenetV3(nn.Module):
def __init__(self):
super().__init__()
self.model = models.mobilenet_v3_small()
self.features = []
self.layers = [3,8,11]
for i,ln in enumerate(self.layers):
self.model.features[ln].register_forward_hook(self.get_features())
def get_features(self):
def hook(model, input, output):
self.features.append(output)
return hook
def forward(self, x):
x = self.model(x)
return x, self.features
This works fine If I do this:
device = torch.device('cpu')
x = torch.randn(1,3,192,192).to(device)
mymodel = CustomMobilenetV3()
o, features = mymodel(x)
for i in features:
print(i.shape)
It does not work If I deepcopy the model, then run it:
deep_model = deepcopy(mymodel)
o, features = deep_model(x)
print(features)
for i in features:
print(i.shape)
in this case the variable features remains empty.
Interestingly, this works fine:
def get_features():
def hook(model, input, output):
feats.append(output)
return hook
model = models.mobilenet_v3_small()
layers = [3,8,11]
for i,ln in enumerate(layers):
model.features[ln].register_forward_hook(get_features())
deep_model = deepcopy(model)
feats = []
o = deep_model(x)
for f in feats:
print(f.shape)
What am I missing?
I need to deepcopy the model to use EMA.
Thanks!