手把手教你在Detectron2中搭建自己的Backbone

手把手教你在Detectron2中搭建自己的Backbone

0、准备主干网络

准备一个Bockbone，我们以最新的RepVGG为例：

RepVGG的链接如下：DingXiaoH/RepVGG: RepVGG: Making VGG-style ConvNets Great Again (github.com)

查看：yaml文件：OUT_FEATURES: [“res2”, “res3”, “res4”, “res5”]

修改主干网络，返回一个字典

class RepVGG(nn.Module):

    def __init__(self, num_blocks, width_multiplier=None, override_groups_map=None, deploy=False, use_se=False):
        super(RepVGG, self).__init__()

        assert len(width_multiplier) == 4

        self.deploy = deploy
        self.override_groups_map = override_groups_map or dict()
        self.use_se = use_se

        assert 0 not in self.override_groups_map

        self.in_planes = min(64, int(64 * width_multiplier[0]))

        self.stage0 = RepVGGBlock(in_channels=3, out_channels=self.in_planes, kernel_size=3, stride=2, padding=1, deploy=self.deploy, use_se=self.use_se)
        self.cur_layer_idx = 1
        self.stage1 = self._make_stage(int(64 * width_multiplier[0]), num_blocks[0], stride=2)
        self.stage2 = self._make_stage(int(128 * width_multiplier[1]), num_blocks[1], stride=2)
        self.stage3 = self._make_stage(int(256 * width_multiplier[2]), num_blocks[2], stride=2)
        self.stage4 = self._make_stage(int(512 * width_multiplier[3]), num_blocks[3], stride=2)
        # self.gap = nn.AdaptiveAvgPool2d(output_size=1)
        # self.linear = nn.Linear(int(512 * width_multiplier[3]), num_classes)


    def _make_stage(self, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        blocks = []
        for stride in strides:
            cur_groups = self.override_groups_map.get(self.cur_layer_idx, 1)
            blocks.append(RepVGGBlock(in_channels=self.in_planes, out_channels=planes, kernel_size=3,
                                      stride=stride, padding=1, groups=cur_groups, deploy=self.deploy, use_se=self.use_se))
            self.in_planes = planes
            self.cur_layer_idx += 1
        return nn.Sequential(*blocks)

    def forward(self, x):
        output = {}
        out = self.stage0(x)
        out = self.stage1(out)
        output["res2"] = out
        out = self.stage2(out)
        output["res3"] = out
        out = self.stage3(out)
        output["res4"] = out
        out = self.stage4(out)
        output["res5"] = out

        # out = self.gap(out)
        # out = out.view(out.size(0), -1)
        # out = self.linear(out)
        return output

1、网络需要继承Backbone基类

class RepVGG(nn.Module):->class RepVGG(Backbone):

2、继承Backbone类的output_shape(self)方法

def __init__(self, num_blocks, width_multiplier=None, override_groups_map=None, deploy=False, use_se=False):
    super(RepVGG, self).__init__()
    ...
    self.out_features_names = ["res2", "res3", "res4", "res5"]
    self._out_feature_channels = { # 每一个stage输出的通道数
        'res2': 48,
        'res3': 96,
        'res4': 192,
        'res5': 1280,
    }
    self._out_feature_strides = { # 步幅
        'res2': 4,
        'res3': 8,
        'res4': 16,
        'res5': 32,
    }

def output_shape(self):
    """
	Returns:
		dict[str->ShapeSpec]
	"""
    # this is a backward-compatible default
    return {
        name: ShapeSpec(
            channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
        )
        for name in self.out_features_names
    }

3、注册RepVGG主干网络

这里这些参数也可以通过cfg来专递，这里简便起见，直接返回模型。

@BACKBONE_REGISTRY.register()
def build_repvgg_backbone(cfg, input_shape):
    deploy = None
    model = RepVGG(num_blocks=[2, 4, 14, 1],
                  width_multiplier=[0.75, 0.75, 0.75, 2.5], override_groups_map=None, deploy=deploy)
    return model

4、在_init_()中写入build_repvgg_backbone

from .repvgg import build_repvgg_backbone

5、删除目录中的build文件，重新编译网络

python setup.py build develop

6、在yaml文件中写入repvgg

MODEL:
    NAME: "build_repvgg_backbone"

7、附录 repvgg预训练模型下载地址

百度网盘 请输入提取码 (baidu.com)

提取码：rvgg

参考链接：DingXiaoH/RepVGG: RepVGG: Making VGG-style ConvNets Great Again (github.com)