detectron2平台注册主干网络

detectron2平台注册主干网络

在detectron2平台训练时，新增的主干网络、检测头需要注册。

---

---

前言

---

一、官方网址

detectron2官方文档链接

二、注册主干网络

1.官网原文如下

Write Models

If you are trying to do something completely new, you may wish to implement a model entirely from scratch. However, in many situations you may be interested in modifying or extending some components of an existing model. Therefore, we also provide mechanisms that let users override the behavior of certain internal components of standard models.

Register New Components

For common concepts that users often want to customize, such as “backbone feature extractor”, “box head”, we provide a registration mechanism for users to inject custom implementation that will be immediately available to use in config files.

For example, to add a new backbone, import this code in your code:

from detectron2.modeling import BACKBONE_REGISTRY, Backbone, ShapeSpec

@BACKBONE_REGISTRY.register()
class ToyBackbone(Backbone):
  def __init__(self, cfg, input_shape):
    super().__init__()
    # create your own backbone
    self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=16, padding=3)

  def forward(self, image):
    return {"conv1": self.conv1(image)}

  def output_shape(self):
    return {"conv1": ShapeSpec(channels=64, stride=16)}

In this code, we implement a new backbone following the interface of the Backbone class, and register it into the BACKBONE_REGISTRY which requires subclasses of Backbone. After importing this code, detectron2 can link the name of the class to its implementation. Therefore you can write the following code:

cfg = ...   # read a config
cfg.MODEL.BACKBONE.NAME = 'ToyBackbone'   # or set it in the config file
model = build_model(cfg)  # it will find `ToyBackbone` defined above

As another example, to add new abilities to the ROI heads in the Generalized R-CNN meta-architecture, you can implement a new ROIHeads subclass and put it in the ROI_HEADS_REGISTRY. DensePose and MeshRCNN are two examples that implement new ROIHeads to perform new tasks. And projects/ contains more examples that implement different architectures.

A complete list of registries can be found in API documentation. You can register components in these registries to customize different parts of a model, or the entire model.

2.按照以上格式注册主干网络

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

python setup.py build develop

2.2 在yaml文件中写入网络，比如repvgg

MODEL:
    NAME: "build_repvgg_backbone"

3.报错

3.1 KeyError: ‘Non-existent config key: MODEL.×××××××ב

.yaml文件新增变量后，报错。
解决方法：
在train_net.py中，setup()函数中，一定要有add_***_config(cfg)

def setup(args):
    """
    Create configs and perform basic setups.
    """
    cfg = get_cfg()
    # 新增这一句，根据自己的网络定义
    add_******_config(cfg)
    cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()
    default_setup(cfg, args)
    return cfg

然后来到add_***_config(）函数代码，添加刚才在ymal配置文件新增的变量。

def add_****_config(cfg):
	cfg.MODEL.temp = True  # 新增加的变量temp

3.2 CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling `cublasCreate(handle)

在Starting training from iteration 0阶段
return torch._C._nn.linear(input, weight, bias) ，报错如上。结合
链接中，对该问题的描述，可能是cpu和gpu都运算时，没有初始化cuda().

添加后重新运行，报错变为

报错变为内存不够，选择合适的gpu在服务器上运行。

三.参考文档

1.手把手教你在Detectron2中搭建自己的Backbone
2.解决bug ： KeyError: ‘Non-existent config key: MODEL.×××××××ב