torchvision中提供了通过Resnet + FPN的方式构建基础网络,这里以resnet50为例实现基础网络的构建。
目录
BackboneWithFPN
类BackboneWithFPN的在backbon_utils.py中,功能是给骨干网络加上FPN,定义如下:
class BackboneWithFPN(nn.Sequential):
def __init__(self, backbone, return_layers, in_channels_list, out_channels):
body = IntermediateLayerGetter(backbone, return_layers=return_layers)
fpn = FeaturePyramidNetwork(
in_channels_list=in_channels_list,
out_channels=out_channels,
extra_blocks=LastLevelMaxPool(),
)
super(BackboneWithFPN, self).__init__(OrderedDict(
[("body", body), ("fpn", fpn)]))
self.out_channels = out_channels
使用时的几个参数:
- backbone : 骨干网络
- return_layers : 需要做fpn的layer的dict
- in_channels_list : return_layer对应的channels
一个简单的使用例子:
import torch
import torchvision.models as models
import torchvision.models.detection.backbone_utils as backbone_utils
backbone = models.resnet50()
return_layers = {'layer1': 0, 'layer2': 1, 'layer3': 2, 'layer4': 3}
in_channels_list = [256,512,1024,2048]
out_channels = 256
resnet_with_fpn = backbone_utils.BackboneWithFPN(backbone,
return_layers,in_channels_list,out_channels)
resnet_with_fpn返回的是一个orderDict,包含了5个fpn后的feature,channel分别为160,80,40,20,10:
input = torch.Tensor(4,3,640,640)
output = resnet_with_fpn(input)
# output是一个OrderedDict
print(type(output))
>>> <class 'collections.OrderedDict'>
for k,v in output.items():
print(k)
print(v)
>>> 0
torch.Size([4, 256, 160, 160])
1
torch.Size([4, 256, 80, 80])
2
torch.Size([4, 256, 40, 40])
3
torch.Size([4, 256, 20, 20])
pool
torch.Size([4, 256, 10, 10])
下面要讲一下BackboneWithFPN中的两个重要组件:
-
IntermediateLayerGetter
-
FeaturePyramidNetwork
IntermediateLayerGetter
关于IntermediateLayerGetter,torchvision给出定义是这样的:
Module wrapper that returns intermediate layers from a model
参数定义:
- model : 输出特征的网络
- return_layers : 一个dict,key是要输出的feature名字,value是要被封装成的新名字。
其实意思就是返回特定层的feature,这个feature通过输入return_layers参数来控制,看一下官方给的使用例子:
Examples::
>>> m = torchvision.models.resnet18(pretrained=True)
>>> # extract layer1 and layer3, giving as names `feat1` and feat2`
>>> new_m = torchvision.models._utils.IntermediateLayerGetter(m,
>>> {'layer1': 'feat1', 'layer3': 'feat2'})
>>> out = new_m(torch.rand(1, 3, 224, 224))
>>> print([(k, v.shape) for k, v in out.items()])
>>> [('feat1', torch.Size([1, 64, 56, 56])),
>>> ('feat2', torch.Size([1, 256, 14, 14]))]
可以看到通过指定参数return_layers={'layer1': 'feat1', 'layer3': 'feat2'},IntermediateLayerGetter返回了resnet18中的layer1和layer3,并封装成了新的名字'feat1'和'feat2',这也符合我们的设定,从feature size上看,layer1做了4倍下采样,layer3做了16倍下采样。
注意:这里的return_layer返回的一定是包含这个key的最后一个layer。比如在resnet中,我们知道layer1是一整个block组成的模块,IntermediateLayerGetter返回的将是整个layer1模块中的最后一层。
FeaturePyramidNetwork
FeaturePyramidNetwork的定义:
Module that adds a FPN from on top of a set of feature maps.
其实这个才是BackboneWithFPN的核心组件,功能也很明确,将FPN加到给定的feature maps中。
参数定义:
- in_channels_list : 输入feature的channels
- out_channels : 输出feature的channels(统一值)
- extra_blocks : 如果你想对某些层做额外的操作需要传入这个参数,注意extra_blocks一定要继承自ExtraFPNBlock类。
官方给的使用例子:
Examples::
>>> m = torchvision.ops.FeaturePyramidNetwork([10, 20, 30], 5)
>>> # get some dummy data
>>> x = OrderedDict()
>>> x['feat0'] = torch.rand(1, 10, 64, 64)
>>> x['feat2'] = torch.rand(1, 20, 16, 16)
>>> x['feat3'] = torch.rand(1, 30, 8, 8)
>>> # compute the FPN on top of x
>>> output = m(x)
>>> print([(k, v.shape) for k, v in output.items()])
>>> # returns
>>> [('feat0', torch.Size([1, 5, 64, 64])),
>>> ('feat2', torch.Size([1, 5, 16, 16])),
>>> ('feat3', torch.Size([1, 5, 8, 8]))]
注意FeaturePyramidNetwork的输入和输出均为OrderedDict,可以看到输入的feat0/feat2/feat3经过FPN后通道数均变成了5。
这里觉得FeaturePyramidNetwork的对FPN的实现写的还蛮简洁精妙的,这里对其中的forward部分做一些注释:
def forward(self, x):
# 输入的x是OrdereDict,将key和value解包成list
names = list(x.keys())
x = list(x.values())
# innner_blocks是一个包含很多个1*1卷积的ModuleList
# inner_block通过1*1卷积改变输入feature的通道数到out_channels
# 对于top的feature只做inner_block输出就好了
last_inner = self.inner_blocks[-1](x[-1])
results = []
results.append(self.layer_blocks[-1](last_inner))
# 因为已经对top feature做过处理,所以循环中[:-1][::-1]刨除对top feature的操作
for feature, inner_block, layer_block in zip(
x[:-1][::-1], self.inner_blocks[:-1][::-1], self.layer_blocks[:-1][::-1]
):
if not inner_block:
continue
# 先对输入的feature的inner_block的1*1卷积
inner_lateral = inner_block(feature)
feat_shape = inner_lateral.shape[-2:]
# last_inner是保存的邻近top层的feature
# interpolate即上采样upsample
inner_top_down = F.interpolate(last_inner, size=feat_shape, mode="nearest")
# 融合当前feature和top层feature
last_inner = inner_lateral + inner_top_down
# 降feature插入到result list的最前面
results.insert(0, layer_block(last_inner))
# 如果定义了extra_blocks,则对x特定的层做额外操作
# 比如这里实现的fpn中对last feature做了maxpooling后当做一个新的feature输出
if self.extra_blocks is not None:
results, names = self.extra_blocks(results, x, names)
# make it back an OrderedDict
out = OrderedDict([(k, v) for k, v in zip(names, results)])
return out
可以对照着下图看一下代码中的变量和注释(灵魂画手只能帮大家到这里了):
至此,torchvision中backbone + fpn的操作实现完毕,不妨来看看torchvision是怎么使用的,在torchvision/models/detection/backbone_utils.py中可以看到这部分代码:
def resnet_fpn_backbone(backbone_name, pretrained):
# 定义backbone(resnet)
backbone = resnet.__dict__[backbone_name](
pretrained=pretrained,
norm_layer=misc_nn_ops.FrozenBatchNorm2d)
# 冻结除layer2/layer3/layer4外的其他层
for name, parameter in backbone.named_parameters():
if 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
parameter.requires_grad_(False)
# 需要做fpn的层是layer0~3,并且将fpn返回的feature封装了新的名字‘0’~‘3’
return_layers = {'layer1': 0, 'layer2': 1, 'layer3': 2, 'layer4': 3}
# 输出特征通道统一为256
in_channels_stage2 = 256
# 对应layer0~3的输入通道,为256~256*8
in_channels_list = [
in_channels_stage2,
in_channels_stage2 * 2,
in_channels_stage2 * 4,
in_channels_stage2 * 8,
]
out_channels = 256
# 返回BackboneWithFPN实例
return BackboneWithFPN(backbone, return_layers, in_channels_list, out_channels)
这样通过调用resnet_fpn_backbone函数,就会生成一个对应resnet+fpn的骨干网络模型。