class ResNet(nn.Module):
def __init__(self, output_features, fine_tuning=False):
super(ResNet, self).__init__()
self.resnet152 = tv.models.resnet152(pretrained=True)
# 冻结特征提取层
for param in self.resnet152.parameters():
param.requires_grad = fine_tuning
#self.features = self.resnet152.features
self.num_fts = 512
self.output_features = output_features
# 从512到1024的线性层
self.classifier = nn.Linear(self.num_fts, self.output_features)
nn.init.xavier_uniform_(self.classifier.weight)
self.tanh = nn.Tanh()
def forward(self, x):
h = self.resnet152(x)
print('h: ',h.shape)
return h
image_model_resnet152=ResNet(output_features=10).to(device)
image_model_resnet152
在这里,打印image_model_resnet152
后,我得到:
这里,(avgpool): Linear(in_features=2048)
和(classifier): Linear(in_features=512)
有什么区别?
我正在实现一个图像标注模型,那么对于图像,我应该选择哪个in_features
?
回答:
ResNet并不像VGG那样简单:它不是一个顺序模型,即在torchvision.models.resnet152
的forward
定义中有一些特定于模型的逻辑,例如,在CNN和分类器之间的特征扁平化。你可以查看其源代码。
在这种情况下,最简单的方法是在CNN的最后一层layer4
上添加一个钩子,并将该层的输出记录在一个外部的dict
中。这是通过register_forward_hook
完成的。
定义钩子:
out = {}
def result(module, input, output):
out['layer4'] = output
将钩子附加到子模块resnet.layer4
上:
>>> x = torch.rand(1,3,224,224)
>>> resnet = torchvision.models.resnet152()
>>> resnet.layer4.register_forward_hook(result)
推理后,你将能够访问out
中的结果:
>>> resnet(x)
>>> out['layer4']
(1, 2048, 7, 7)
你可以查看我关于前向钩子更深入使用的另一个回答。
一个可能的实现方式是:
class NN(nn.Module):
def __init__(self):
super().__init__()
self.resnet = torchvision.models.resnet152()
self.resnet.layer4.register_forward_hook(result)
self.out = {}
@staticmethod
def result(module, input, output):
out['layer4'] = output
def forward(self, x):
x = self.resnet(x)
return out['layer4']
然后你可以为你的自定义分类器定义额外的层,并在forward
中调用它们。