【Pytorch】调用中间层的结果

在研究 Retinaface 网络结构时候，有个疑惑，作者怎么把 MobileNetV1 的三个 stage 输出分别接到 FPN 上面的，我注意到下面的代码：

import torchvision.models._utils as _utils
# 使用 _utils.IntermediateLayerGetter 函数获得中间层的结果，第一个参数时网络，第二个参数时字典
self.body = _utils.IntermediateLayerGetter(backbone, cfg['return_layers'])

test

class MobileNetV1(nn.Module):
	def __init__(self):
		super(MobileNetV1, self).__init__()
		self.stage1 = nn.Sequential(
			conv_bn(3, 8, 2, leaky = 0.1),    # 3
			conv_dw(8, 16, 1),   # 7
			conv_dw(16, 32, 2),  # 11
			conv_dw(32, 32, 1),  # 19
			conv_dw(32, 64, 2),  # 27
			conv_dw(64, 64, 1),  # 43
		)
		self.stage2 = nn.Sequential(
			conv_dw(64, 128, 2),  # 43 + 16 = 59
			conv_dw(128, 128, 1), # 59 + 32 = 91
			conv_dw(128, 128, 1), # 91 + 32 = 123
			conv_dw(128, 128, 1), # 123 + 32 = 155
			conv_dw(128, 128, 1), # 155 + 32 = 187
			conv_dw(128, 128, 1), # 187 + 32 = 219
		)
		self.stage3 = nn.Sequential(
			conv_dw(128, 256, 2), # 219 +3 2 = 241
			conv_dw(256, 256, 1), # 241 + 64 = 301
		)
		self.avg = nn.AdaptiveAvgPool2d((1,1))
		self.fc = nn.Linear(256, 1000)
 
	def forward(self, x):
		x = self.stage1(x)
		x = self.stage2(x)
		x = self.stage3(x)
		x = self.avg(x)
		# x = self.model(x)
		x = x.view(-1, 256)
		x = self.fc(x)
		return x

'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3},