Detectron2: Code Notes

Posted on Edited on

Detectron 2 源码阅读笔记

Most part of the begining of this notes is inherited from https://www.cnblogs.com/marsggbo/p/11677086.html.

1. 代码结构概览

  1. 核心部分

    • configs: 储存各种网络的yaml配置文件；
    • datasets: 存放数据集的地方；
    • deteectron2：运行代码的核心组建；
    • tools：提供了运行代码的入口以及一切可视化的代码问价。

  2. Tutorial部分

    • demo：demo；
    • docs：docs；
    • tests：tests；
    • projects：提供了真实的项目代码示例；

2. 代码逻辑分析

2.1 超参数配置

Detectron2中的参数配置使用了yacs这个库，这个库能够很好地重用和拼接超参数文件配置。我们先看一下detrctron2/config/的文件结构：

  • compat.py: 应该是对之前的Detectron库的兼容吧，可忽略。
  • config.py: 定义了一个CfgNode类，这个类继承自fvcore库(fb写的一个共公共库，提供一些共享的函数，方便各种不同项目使用)中定义的CfgNode,总之就是不断继承。。。继承关系是这样的：detrctron2.config.CfgNode->fcvore.common.config.CfgNode->yacs.config.CfgNode->dict。另外该文件还提供了get_cfg()方法，该方法会返回一个含有默认配置的CfgNode,而这些默认的配置值在下面的default.py中定义了，之所以这样做是因为要配置的默认值太多了，所以为了文档清晰才写到了一个新的文件中去，不过，yacs库的作者也建议这样做。
  • default.py: 如上面所说，该文件定义了各种参数的默认值。

了解配置函数的方法后我们再回到tools/train_net.py，我们一行一行的来理解。

  • tools/train_net.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
from detectron2.config import get_cfg
from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, hooks, launch
...

def setup(args):
    """
    Create configs and perform basic setups.
    """
    cfg = get_cfg() 
    cfg.merge_from_file(args.config_file) 
    cfg.merge_from_list(args.opts)
    cfg.freeze()
    default_setup(cfg, args)
    return cfg
  • cfg = get_cfg()： 获取已经配置好默认参数的cfg
  • cfg.merge_from_file(args.config_file)：config_file是指定的yaml配置文件，通过merge_from_file这个函数会将yaml文件中指定的超参数对默认值进行覆盖。
  • cfg.merge_from_list(args.opts)：merge_from_list作用同上面的类似，只不过是通过命令行的方式覆盖。
  • cfg.freeze(): freeze函数的作用是将超参数值冻结，避免被程序不小心修改。
  • default_setup(cfg, args)：default_setup是detectron2/engine/default.py中提供的一个默认配置函数，具体是怎么配置的这里不详细说明了。不过需要知道的值这个文件中还提供了很多其他的配置函数，例如还提供了两个类：DefaultPredictor和DefaultTrainer。

2.2 Trainer结构

既然上面提到了DefaultTrainer，那么我们就从这个类入手了解一下detectron2.engine,其代码结构如下：

  • train_loop.py：这个函数主要作用是提供了三个重要的类：

    • HookBase: 这是一个Hook的基类，用于指定在训练前后或者每一个step前后需要做什么事情，所以根据特定的需求需要对如下四种方法做不同的定义：before_train, after_train, before_step, after_step。以before_step。
    • TrainerBase: 该类中定义的函数可以归纳成三种，并初始化_hooks为[]：
      • register_hooks:这个很好理解，就是将用户定义的一些hooks进行注册，说大白话就是把若干个Hook放在一个list里面去。之后只需要遍历这个list依次执行就可以了。
        • register_hooks(hooks)会将输入的hooks中的trainer设置为当前的类实例的弱引用，并将其添加到当前类实例的self._hooks中去；
      • 第二类其实就是上面提到的遍历hook list并执行hook，不过这个遍历有四种，分别是before_train,after_train,before_step,after_step。还有一个就是run_step,这个函数其实就是平常我们在编写训练过程的代码，例如读数据，训练模型，获取损失值，求导数，反向梯度更新等,只不过在这个类里面没有定义。
      • 第三类就是train函数，它有两个参数，分别是开始的迭代数和最大的迭代数。之后就是重复依次执行第二类中的函数指定迭代次数。
    • SimpleTrainer:其实就是继承自TrainerBase,然后定义了run_step等方法。我们后面也可以继承这个类做进一步的自定义。
      • 初始化基类，并赋予其model、data_loader、_data_loader_iter、optimizer；
      • 这在随后的训练过程中会被继承为默认的训练类DefaultTrainer。

  • defaults.py: 上面已介绍，提供了两个类：DefaultPredictor和DefaultTrainer，这个DefaultTrainer就继承自SimpleTrainer,所以存在如下继承关系：detectron2.engine.default.DefaultTrainer->detectron2.engine.train_loop.SimpleTrainer->detectron2.engine.train_loop.TrainerBase

  • hooks.py:定义了很多继承自train_loop.HookBase的Hook。

    • CallBackHook: Create a hook using callback functions provided by the user.
    • IterationTimer：Track the time spent for each iteration (each run_step call in the trainer). Print a summary in the end of training.；
    • PeriodicWriter：Write events to EventStorage periodically.
    • PeriodicCheckpointer：Same as :class:detectron2.checkpoint.PeriodicCheckpointer, but as a hook.
    • LRScheduler：A hook which executes a torch builtin LR scheduler and summarizes the LR. It is executed after every iteration.
    • AutogradProfiler：A hook which runs torch.autograd.profiler.profile.
    • EvalHook: Run an evaluation function periodically, and at the end of training.
    • PreciseBN: The standard implementation of BatchNorm uses EMA in inference, which is sometimes suboptimal. This class computes the true average of statistics rather than the moving average, and put true averages to every BN layer in the given model.

  • launch.py: 前面提到过，可以理解成代码启动器，可以根据命令决定是否采用分布式训练（或者单机多卡）或者单机单卡训练。

整个训练过程按照如下流程进行：

1
2
3
4
5
6
7
self.before_train()
for self.iter in range(start_iter, max_iter):
  self.before_step()
  self.run_step()
  self.after_step()
finally:
  self.after_train()

小结：
fig-1

2.3 Trainer解析

1
2
3
4
5
6
7
8
9
10
11
12
def main(args):
    cfg = setup(args)

    if args.eval_only:
        ...
    trainer = Trainer(cfg)
    trainer.resume_or_load(resume=args.resume)
    if cfg.TEST.AUG.ENABLED:
        trainer.register_hooks(
            [hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]
        )
    return trainer.train()

2.3.1 build_*方法

build_*主要来自于engine.defaults.DefaultTrainer类，继承自engine.train_loop.SimpleTrainer类。而detectron2.engine.default.DefaultTrainer在其__init__(self, cfg)函数中定义了解析cfg。如下面代码所示，cfg会作为参数倍若干个build_*方法解析，得到解析后的model,optimizer,data_loader等。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
from detectron2.modeling import build_model
class DefaultTrainer(SimpleTrainer):
  def __init__(self, cfg):
    """
    Args:
        cfg (CfgNode):
    """
    # Assume these objects must be constructed in this order.
    model = self.build_model(cfg)
    optimizer = self.build_optimizer(cfg, model)
    data_loader = self.build_train_loader(cfg)
    
    ... 
    
    self.register_hooks(self.build_hooks())
      
  @classmethod
  def build_model(cls, cfg):
    """
    Returns:
        torch.nn.Module:
    """
    model = build_model(cfg)
    logger = logging.getLogger(__name__)
    logger.info("Model:\n{}".format(model))
    return model

下面我们以DefaultTrainer.build_model为例来介绍注册机制,该方法调用了detectron2/modeling/meta_arch/build.py的build_model函数,其源代码如下：

1
2
3
4
5
6
7
8
9
10
11
from detectron2.utils.registry import Registry

META_ARCH_REGISTRY = Registry("META_ARCH")
META_ARCH_REGISTRY.__doc__ = """..."""

def build_model(cfg):
  """
  Built the whole model, defined by `cfg.MODEL.META_ARCHITECTURE`.
  """
  meta_arch = cfg.MODEL.META_ARCHITECTURE
  return META_ARCH_REGISTRY.get(meta_arch)(cfg)
  • meta_arch = cfg.MODEL.META_ARCHITECTURE： 根据超参数获得网络结构的名字
  • return META_ARCH_REGISTRY.get(meta_arch)(cfg)：META_ARCH_REGISTRY是一个Registry类(这个在后面会详细介绍)，可以将这一行代码拆成如下几个步骤：
1
2
model = META_ARCH_REGISTRY.get(meta_arch)
return model(cfg)

2.3.2 注册机制Registry

假如你想自己实现一个新的backbone网络，那么你可以这样做:

首先在detectron2中定义好如下（实际上已经定义了）：

1
2
3
# detectron2/modeling/backbone/build.py

BACKBONE_REGISTRY = Registry('BACKBONE')

之后在你创建的新的文件下按如下方式创建你的backbone

1
2
3
4
5
6
7
8
9
10
11
12
# detectron2/modeling/backbone/your_backbone.py
from .build import BACKBONE_REGISTRY

# 方式1
@BACKBONE_REGISTRY.register()
class MyBackbone():
    ...
        
# 方式2
class MyBackbone():
    ...
BACKBONE_REGISTRY.register(MyBackbone)

Registry源代码如下（有删减）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class Registry(object):
  def __init__(self, name):
    self._name = name
    self._obj_map = {}

  def _do_register(self, name, obj):
    assert (
      name not in self._obj_map
    ), "An object named '{}' was already registered in '{}' registry!".format(name, self._name)
    self._obj_map[name] = obj

  def register(self, obj=None):
    if obj is None:
      # used as a decorator
      def deco(func_or_class):
        name = func_or_class.__name__
        self._do_register(name, func_or_class)
        return func_or_class

      return deco

    # used as a function call
    name = obj.__name__
    self._do_register(name, obj)

  def get(self, name):
    ret = self._obj_map.get(name)
    if ret is None:
      raise KeyError("No object named '{}' found in '{}' registry!".format(name, self._name))
    return ret
  • 首先是__init__部分：
    • self._name则是你要注册的名字，例如对于完整的模型而言，name一般取META_ARCH。当然如果你需要自定义backbone网络，你也可以定义一个Registry('BACKBONE');
    • self._obj_map:其实就是一个字典。以模型为例，key就是你的模型名字，而value就是对应的模型类。这样你在传参时只需要修改一下模型名字就能使用不同的模型了。具体实现方法就是后面这几个函数;
  • register: 可以看到该方法定义了注册的两种方式，一种是当obj==None的时候，使用装饰器的方式注册，另外一种就是直接将obj作为参数调用_do_register进行注册。
  • _do_register:真正注册的函数，可以看到它首先会判断name是否已经存在于self._obj_map了。什么意思呢？还是以backbone为例，我们定义了一个BACKBONE_REGISTRY = Registry('BACKBONE'),然后又定义了很多种backbone，而这些backbone都使用@BACKBONE_REGISTRY.register()的方式注册到了BACKBONE_REGISTRY._obj_map中了,所以才取名为Registry。

2.4 Detectron2 整体代码框架

整体框架如下图所示：

fig-2

3. Dataset

构建data_loader原理步骤

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# engine/default.py
from detectron2.data import (
  MetadataCatalog,
  build_detection_test_loader,
  build_detection_train_loader,
)

class DefaultTrainer(SimpleTrainer):
  def __init__(self, cfg):
    # Assume these objects must be constructed in this order.
    data_loader = self.build_train_loader(cfg)
  
    ...
  
  @classmethod
  def build_train_loader(cls, cfg):
    """
    Returns:
      iterable
    """
    return build_detection_train_loader(cfg)

函数调用关系如下图：

fig-3

结合前面两篇文章的内容可以看到detectron2在构建model,optimizer和data_loader的时候都是在对应的build.py文件里实现的。我们看一下build_detection_train_loader是如何定义的(对应上图中紫色方框内的部分(自下往上的顺序)）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def build_detection_train_loader(cfg, mapper=None):
  """
  A data loader is created by the following steps:

  1. Use the dataset names in config to query :class:`DatasetCatalog`, and obtain a list of dicts.
  2. Start workers to work on the dicts. Each worker will:
    * Map each metadata dict into another format to be consumed by the model.
    * Batch them by simply putting dicts into a list.
  The batched ``list[mapped_dict]`` is what this dataloader will return.

  Args:
    cfg (CfgNode): the config
    mapper (callable): a callable which takes a sample (dict) from dataset and
      returns the format to be consumed by the model.
      By default it will be `DatasetMapper(cfg, True)`.

  Returns:
    a torch DataLoader object
  """
# 1. 获得dataset_dicts
  dataset_dicts = get_detection_dataset_dicts(
    cfg.DATASETS.TRAIN,
    filter_empty=True,
    min_keypoints=cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
    if cfg.MODEL.KEYPOINT_ON
    else 0,
    proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
  )

# 2. 将dataset_dicts转化成torch.utils.data.Dataset
  dataset = DatasetFromList(dataset_dicts, copy=False)

# 3. 进一步转化成MapDataset，每次读取数据时都会调用mapper来对dict进行解析
  if mapper is None:
    mapper = DatasetMapper(cfg, True)
  dataset = MapDataset(dataset, mapper)

# 4. 采样器
  sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
  if sampler_name == "TrainingSampler":
    sampler = samplers.TrainingSampler(len(dataset))

  ...

  batch_sampler = build_batch_data_sampler(
    sampler, images_per_worker, group_bin_edges, aspect_ratios
  )

# 5. 数据迭代器 data_loader
  data_loader = torch.utils.data.DataLoader(
    dataset,
    num_workers=cfg.DATALOADER.NUM_WORKERS,
    batch_sampler=batch_sampler,
    collate_fn=trivial_batch_collator,
    worker_init_fn=worker_init_reset_seed,
  )

  return data_loader

后面的采样器和data_loader可以参阅https://www.cnblogs.com/marsggbo/p/11308889.html一文弄懂Pytorch的DataLoader, DataSet, Sampler之间的关系。

3.1 获取dataset_dicts

get_detection_dataset_dicts(dataset_names)函数需要传递的一个重要参数是dataset_names,这个参数其实就是一个字符串，用来指定数据集的名称。通过这个字符串，该函数会调用data/catalog.py的DatasetCatalog类来进行解析得到一个包含数据信息的字典。

解析的原理是：DatasetCatalog有一个字典_REGISTERED，默认已经注册好了例如coco, voc这些数据集的信息。如果你想要使用你自己的数据集，那么你需要在最开始前你需要定义你的数据集名字以及定义一个函数(这个函数不需要传参，而且最后会返回一个dict，该dict包含你的数据集信息)，举个栗子：

1
2
3
4
5
6
7
from detectron2.data import DatasetCatalog
my_dataset_name = 'apple'
def get_dicts():
	...
	return dict

DatasetCatalog.register(my_dataset_name, get_dicts)

当然，如果你的数据集已经是COCO的格式了，那么你也可以使用如下方法进行注册：

1
2
3
from detectron2.data.datasets import register_coco_instances
my_dataset_name = 'apple'
register_coco_instances(my_dataset_name, {}, "json_annotation.json", "path/to/image/dir")

最后，get_detection_dataset_dicts会返回一个包含若干个dict的list，之所以是list是因为参数dataset_names也是一个list，这样我们就可以制定多个names来同时对数据进行读取。

3.2 解析成DatasetFromList

DatasetFromList(dataset_dict)函数定义在detectron2/data/common.py中，它其实就是一个torch.utils.data.Dataset类，其源码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class DatasetFromList(data.Dataset):
  """
  Wrap a list to a torch Dataset. It produces elements of the list as data.
  """

  def __init__(self, lst: list, copy: bool = True):
    """
    Args:
      lst (list): a list which contains elements to produce.
      copy (bool): whether to deepcopy the element when producing it,
        so that the result can be modified in place without affecting the
        source in the list.
    """
    self._lst = lst
    self._copy = copy

  def __len__(self):
    return len(self._lst)

  def __getitem__(self, idx):
    if self._copy:
      return copy.deepcopy(self._lst[idx])
    else:
      return self._lst[idx]

3.3 将DatasetFromList转化为MapDataset

其实DatsetFromList和MapDataset都是torch.utils.data.Dataset的子类，那他们的区别是什么呢？很简单，区别就是后者使用了mapper。

在解释mapper是什么之前我们首先要知道的是，在detectron2中，一张图片对应的是一个dict，那么整个数据集就是listdict。之后我们再看DatsetFromList，它的__getitem__函数非常简单，它只是简单粗暴地就返回了指定idx的元素。显然这样是不行的，因为在把数据扔给模型训练之前我们肯定还要对数据做一定的处理，而这个工作就是由mapper来做的，默认情况下使用的是detectron2/data/dataset_mapper.py中定义的DatasetMapper，如果你需要自定义一个mapper也可以参考这个写。

1
DatasetMapper(cfg, is_train=True)

我们继续了解一下DatasetMapper的实现原理，首先看一下官方给的定义：

1
2
3
4
"""
A callable which takes a dataset dict in Detectron2 Dataset format,
and map it into a format used by the model.
"""

简单概括就是这个类是可调用的(callable),所以在下面的源码中可以看到定义了__call__方法。

该类主要做了这三件事：

1
2
3
4
The callable currently does the following:
1. Read the image from "file_name"
2. Applies cropping/geometric transforms to the image and annotations
3. Prepare data and annotations to Tensor and :class:`Instances`

其源码如下（有删减）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class DatasetMapper:
  def __init__(self, cfg, is_train=True):
  # 读取cfg的参数
  ...

  def __call__(self, dataset_dict):
    """
    Args:
        dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.

    Returns:
        dict: a format that builtin models in detectron2 accept
    """
    dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
  
  # 1. 读取图像数据
    image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
  
  # 2. 对image和box等做Transformation
    if "annotations" not in dataset_dict:
      image, transforms = T.apply_transform_gens(
        ([self.crop_gen] if self.crop_gen else []) + self.tfm_gens, image
      )
    else:
    ...
      image, transforms = T.apply_transform_gens(self.tfm_gens, image)
      if self.crop_gen:
        transforms = crop_tfm + transforms

      image_shape = image.shape[:2]  # h, w
  
  # 3.将数据转化成tensor格式
    dataset_dict["image"] = torch.as_tensor(image.transpose(2, 0, 1).astype("float32"))
  ...

    return dataset_dict

MapDataset:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class MapDataset(data.Dataset):
  def __init__(self, dataset, map_func):
    self._dataset = dataset
    self._map_func = PicklableWrapper(map_func)  # wrap so that a lambda will work

    self._rng = random.Random(42)
    self._fallback_candidates = set(range(len(dataset)))

  def __len__(self):
    return len(self._dataset)

  def __getitem__(self, idx):
    retry_count = 0
    cur_idx = int(idx)

    while True:
      data = self._map_func(self._dataset[cur_idx])
      if data is not None:
        self._fallback_candidates.add(cur_idx)
        return data

        # _map_func fails for this idx, use a random new index from the pool
        retry_count += 1
        self._fallback_candidates.discard(cur_idx)
        cur_idx = self._rng.sample(self._fallback_candidates, k=1)[0]

        if retry_count >= 3:
          logger = logging.getLogger(__name__)
          logger.warning(
            "Failed to apply `_map_func` for idx: {}, retry count: {}".format(
              idx, retry_count
            )
          )
  • self._fallback_candidates是一个set,它的特点是其中的元素是独一无二的，定义这个的作用是记录可正常读取的数据索引，因为有的数据可能无法正常读取，所以这个时候我们就可以把这个坏数据的索引从_fallback_candidates中剔除，并随机采样一个索引来读取数据。
  • __getitem__中的逻辑就是首先读取指定索引的数据，如果正常读取就把该所索引值加入到_fallback_candidates中去；反之，如果数据无法读取，则将对应索引值删除，并随机采样一个数据，并且尝试3次，若3次后都无法正常读取数据，则报错，但是好像也没有退出程序，而是继续读数据，可能是以为总有能正常读取的数据吧;

3. Debug

3.1 Using ground.py

4. Training

4.1 在COCO数据集上训练

Using 1 image per batch, cause the desktop can’t afford 2 images per batch.

1
2
3
python tools/train_net.py \
	--config-file configs/COCO-Keypoints/keypoint_rcnn_R_50_FPN_1x.yaml \
	SOLVER.IMS_PER_BATCH 1 SOLVER.BASE_LR 0.0025

Training parameters:

  • max_mem: 2638M;
  • 89999 iterations;
  • The training log is saved to log_1127_first_training.txt;

4.2 训练流程

主要训练函数为tools/train_net.py中的main函数，launch函数将训练过程分布化，在当前阶段暂时不予深究。

5. Validation

python demo/demo.py
–config-file configs/COCO-Keypoints/keypoint_rcnn_R_50_FPN_1x.yaml
–input comunity/1574912817253.png
–opts MODEL.WEIGHTS weights/COCO-Keypoints/model_final_04e291.pkl