从源代码开始 Detectron2学习笔记(二)
训练自己的数据集(Train on custom dataset)
在上次的学习笔记中,我们重点关注了模型的配置和构建,以及使用COCO预训练模型进行检测并可视化结果。这次学习要尝试的是:从零开始构建自己的数据集,并且利用detectron2进行数据集的注册,模型的训练和评估,以及最终的检测效果展示。虽说在detectron2中,上述的操作仅仅需要几行代码即可完成,但我们仍不能局限于此,还是和上一章一样,从源代码入手进行详解。由于这两个月都在忙着搞毕设写论文、一直没能完成第二章的笔记,在此对关注的同学们说声抱歉。
1.数据集的处理与注册(Download&Registry)
1.1 数据的下载与处理
这次选用的数据集为只含一个类别的气球数据集,任务为气球目标的检测和分割。数据集的下载和初步处理代码如下(源代码见官方colab教程)
# 下载数据集
# 如果是在linux命令行输入 将命令最前端的!去掉即可
!wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip
!unzip balloon_dataset.zip > /dev/null
#数据集的处理
from detectron2.structures import BoxMode
from detectron2.data import MetadataCatalog, DatasetCatalog
def get_balloon_dicts(img_dir):
# 读取json文件
json_file = os.path.join(img_dir, "via_region_data.json")
with open(json_file) as f:
imgs_anns = json.load(f)
dataset_dicts = []
# 获取图片信息
for idx, v in enumerate(imgs_anns.values()):
record = {}
filename = os.path.join(img_dir, v["filename"])
height, width = cv2.imread(filename).shape[:2]
record["file_name"] = filename
record["image_id"] = idx
record["height"] = height
record["width"] = width
# 获取目标的标注信息
annos = v["regions"]
objs = []
for _, anno in annos.items():
assert not anno["region_attributes"]
anno = anno["shape_attributes"]
px = anno["all_points_x"]
py = anno["all_points_y"]
poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
poly = [p for x in poly for p in x]
obj = {
"bbox": [np.min(px), np.min(py), np.max(px), np.max(py)],
"bbox_mode": BoxMode.XYXY_ABS,
"segmentation": [poly],
"category_id": 0,
}
objs.append(obj)
# 返回的dataset_dicts 是由一个个record组成的列表
record["annotations"] = objs
dataset_dicts.append(record)
return dataset_dicts
这段代码并不复杂,定义的函数会读取气球数据集并返回一个由很多个"record"组成的列表。其中,每一个record都代表一张图片及其标注。数据集的结构展示如下:
# 数据集结构:
dataset_dicts
---record 1
---filename
---image_id
---weight
---height
---annotations
---obj 1
---bbox # 检测bounding box标注
---bbox_mode # bounding box 类别 这里是(xmin,ymin,xmax,ymax)
---segmentation # 分割标注
---category_id # 类别信息,因为数据集中只有一个类别,所以此处均为0
--obj 2
--obj 3
...
--record 2
--record 3
...
1.2 在detectron2中进行数据注册
1.2.1 Datasetcatelog
for d in ["train", "val"]:
DatasetCatalog.register("balloon_" + d, lambda d=d: get_balloon_dicts("balloon/" + d))
MetadataCatalog.get("balloon_" + d).set(thing_classes=["balloon"])
balloon_metadata = MetadataCatalog.get("balloon_train")
'''
可能有小伙伴不太理解register()参数中lambda d=d: get_balloon_dicts("balloon/" + d)的意思,在这里稍微举例解释一下,大家也可以在自己的python编译器中尝试:
例如,f = lambda x:x+1 定义了一个隐式的函数,在调用f时必须加入参数如f(1),f(2),否则会报错;
但是,就像python中define定义的函数有默认参数一样,lambda定义的函数也可以有默认参数,如:
f = lambda x=1: x+1 其实等价于:define f(x=1): return x+1, 此时使用f()或f(1)均是正确的用法。
'''
此处是今天的第一个重点知识:在detectron2中,我们的数据是如何组织并且最终与模型对接的呢?首先来看一下跟DatasetCatalog相关的源码detectron2/data/catalog.py
# line 81
DatasetCatalog = _DatasetCatalog()
可以看出,前面的代码中用到的“DatasetCatalog”就是在这个文件中创建的属于 "_Dataset-catalog"类的对象。"Catalog"的中文意思是目录,编目,而个人觉得在detectron2中更适合的翻译是注册表,或者登记表。也就是说,所谓数据注册,就是利用这个注册表,将自己的数据在detectron2中进行登记。
实际上,不只是数据集,检测模型的各个结构,比如Backbone, RPN 等都有自己的注册表,我们可以通过注册表来登记自己设计的模块,这在后面的学习中会慢慢涉及。
回到今天的重点,来看一下“_DatasetCatalog”类的源码:
class _DatasetCatalog(UserDict):
"""
一个储存了数据集信息及其使用方式的全局字典。
它包含了一个从字符串(通常是数据集的名称,例如"coco_2014_train")到一个函数的映射,
而这个函数的功能为:解析数据集,并返回list[dict]格式的数据样本。
返回的字典应该遵从Detectron2 数据集格式的要求(详见DATASETS.md)
**如果和detectron2中其他的数据功能一起使用,那么这个注册表可以让我们只使用配置文件中的字符串
就可以更方便地选取不同类型的数据集(本次学习暂不涉及)
下面的函数注释简单易懂,就不全部翻译了,只挑重点翻。
"""
# 前面代码中用到的注册函数,最基本的功能
def register(self, name, func):
"""
Args:
name (str): the name that identifies a dataset, e.g. "coco_2014_train".
func (callable): a callable which takes no arguments and returns a list of dicts.
如果这个func被调用多次,必须保证每次会返回相同的结果.
"""
# 验证func是否可调用
assert callable(func), "You must register a function with `DatasetCatalog.register`!"
# 检查数据集是否已经被注册
assert name not in self, "Dataset '{}' is already registered!".format(name)
# 建立一个从数据集名称(name)到 数据返回函数(func)的映射
self[name] = func
def get(self, name):
"""
Call the registered function and return its results.
Args:
name (str): the name that identifies a dataset, e.g. "coco_2014_train".
Returns:
list[dict]: dataset annotations.
"""
try:
# 根据前面确定的映射得到数据函数
f = self[name]
# 如果数据集没有被注册,则会抛出异常
except KeyError as e:
raise KeyError(
"Dataset '{}' is not registered! Available datasets are: {}".format(
name, ", ".join(list(self.keys()))
)
) from e
# 调用func,并返回得到的数据
return f()
# 返回包含所有数据集名称的列表
def list(self) -> List[str]:
"""
List all registered datasets.
Returns:
list[str]
"""
return list(self.keys())
# 删除某一数据集
def remove(self, name):
"""
Alias of ``pop``.
"""
self.pop(name)
def __str__(self):
return "DatasetCatalog(registered datasets: {})".format(", ".join(self.keys()))
__repr__ = __str__
源代码中,register函数是重点与难点,我们来逐步分析:
首先是函数的输入:
def register(self, name, func):
'''
name (str): the name that identifies a dataset, e.g. "coco_2014_train".
func (callable): a callable which takes no arguments and returns a list of dicts.
如果这个func被调用多次,必须保证每次会返回相同的结果
'''
第一个输入是一个字符串,代表数据集的名称;第二个输入为一个可调用函数,用于返回数据集的样本。在这里我们结合前面的注册气球数据集的代码来分析:
for d in ["train", "val"]:
DatasetCatalog.register("balloon_" + d, lambda d=d: get_balloon_dicts("balloon/" + d))
'''
这里分别注册了训练集(train)和测试集(test),我们只以train为例分析:
两个输入分别为:“balloon_train”,代表训练集名称;第二个输入是一个由lambda定义的函数,
当注册训练集时,register函数的第二个输入为get_balloon_dicts("balloon_train")
'''
可以看到,当注册训练集时,register函数接收了函数get_balloon_dicts(“ballooon_train”)做为输入,简单回顾一下get_balloon_dicts()函数的输出:
def get_balloon_dicts(img_dir):
...
...
# 返回的dataset_dicts 是由一个个record组成的列表
record["annotations"] = objs
dataset_dicts.append(record)
return dataset_dicts
也就是说,Datasetcatalog本身并没有通用的数据读取功能,而是需要我们自己定义自己数据集的读取方式,要有一个按照Detectron2要求的格式返回数据的函数,并将这个函数作为注册时的输入。简而言之,我们要告诉detectron2的注册表:
a.数据集名称
b.数据的获取方式
继续回到register函数中:
def register(self,name,func):
...
assert callable(func), "You must register a function with `DatasetCatalog.register`!"
# 检查数据集是否已经被注册
assert name not in self, "Dataset '{}' is already registered!".format(name)
# 建立一个从数据集名称(name)到 数据返回函数(func)的映射
self[name] = func
后面的代码简洁易懂,也写了中文注释,就不多做赘述了。
紧接着是get函数:
def get(self, name):
try:
# 根据前面确定的映射得到数据函数
f = self[name]
# 如果数据集没有被注册,则会抛出异常
except KeyError as e:
raise KeyError(
"Dataset '{}' is not registered! Available datasets are: {}".format(
name, ", ".join(list(self.keys()))
)
) from e
# 调用func,并返回得到的数据
return f()
get函数只接受一个函数:name,也就是数据集的名称,并返回一个“f()”,注意这里是“f()”而不是“f”,意思是返回func的调用结果。函数看起来有些复杂,是因为用了一个try-except的结构来判断此名称的数据集是否已经被注册。
总结一下,get函数返回已注册函数func的调用结果,也就是字面意义上“get(获取)”到数据,因此要求此数据集必须已经被注册。
1.2.2 Metadatacatalog
讲完了Datasetcatalog,来讲一下另外一个注册表,也就是Metadatacatalog,其源代码文件与datasetcatalog相同。
# line 230
MetadataCatalog = _MetadataCatalog()
与datasetcatalog相似,MetadataCatalog也是类_MetadataCatalog的一个实例。
但在此之前,我们需要先看一下源代码中的一个类:Metadata(这里只展示一些重要的函数,完整的源代码请参考前述链接)
class Metadata(types.SimpleNamespace):
"""
一个支持简单的属性设置/获取功能的类,其功能是保存一个数据集的“metadata”,
并且让它能够被作为全局变量来使用。
例子:
::
# 在某处你想要加载数据的地方:
MetadataCatalog.get("mydataset").thing_classes = ["person", "dog"]
# 在某处你想要打印数据集信息或做可视化:
classes = MetadataCatalog.get("mydataset").thing_classes
"""
....
....
def set(self, **kwargs):
"""
Set multiple metadata with kwargs.
"""
for k, v in kwargs.items():
setattr(self, k, v)
return self
# 根据键获取属性
def get(self, key, default=None):
"""
Access an attribute and return its value if exists.
Otherwise return default.
"""
try:
return getattr(self, key)
except AttributeError:
return default
比较重要的函数就两个,get()和set(),看字面意思也能明白这俩的作用,一个是获取,一个是写入。看完了metadata,再看最终的MetadataCatalog:
class _MetadataCatalog(UserDict):
"""
MetadataCatalog 是一个全局字典,它提供给定数据集的"Metadata"的访问权限.
一个特定名称的metadata是单一的,一旦被创建,它就会一直保持存在,并且可以通过
调用“get(name)”函数来获取,name参数为数据集的名称。
它类似于全局变量,所以不要滥用。
它是为了保存在各个任务中共享并保持一致的知识,比如COCO数据集的类别名称。
"""
def get(self, name):
"""
Args:
name (str): name of a dataset (e.g. coco_2014_train).
Returns:
返回该类别的Metadata实例。如果该实例不存在,就创建一个。
"""
assert len(name)
r = super().get(name, None)
if r is None:
r = self[name] = Metadata(name=name)
return r
def list(self):
"""
将所有注册的“Metadata”以列表形式返回
"""
return list(self.keys())
def remove(self, name):
"""
删除一个metadata
"""
self.pop(name)
def __str__(self):
return "MetadataCatalog(registered metadata: {})".format(", ".join(self.keys()))
1.2.3 Catalog小结
总结一下,在detectron2中,关于数据的注册表有两个,分别是DatasetCatalog和MetadataCatalog。其中,DatasetCatalog用于一个特定任务所需的特定数据集的注册,注册时必须的输入为<数据集名称,获取数据集的可调用函数>;MetadataCatalog用于储存通用的“元数据”,如COCO各个数据集都能共用的类别名称等。
最后,简要回顾一下气球数据集用到的两个数据集注册表:
for d in ["train", "val"]:
DatasetCatalog.register("balloon_" + d, lambda d=d: get_balloon_dicts("balloon/" + d))
MetadataCatalog.get("balloon_" + d).set(thing_classes=["balloon"])
'''
MetadataCatalog.get("balloon_" + d)会创建一个气球数据集的Metadata并返回。set函数则将metadata的“thing_class”属性设定为“[balloon]”,即只有一个类别.
'''
2. 模型导入与训练(Train)
讲完了数据集的加载和注册,再来看看模型的导入与训练。先上tutorial的代码:
from detectron2 import model_zoo
from detectron2.config import get_cfg
from detectron2.utils.logger import setup_logger
# 获取config文件
cfg = get_cfg()
# 从detection2的model_zoo中导入Mask R-CNN + resnet50 + FPN的配置文件
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
# 指定训练数据集为 “balloon_train”
cfg.DATASETS.TRAIN = ("balloon_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
# 指定模型所用的初始化权重
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
# Batch_size = 2
cfg.SOLVER.IMS_PER_BATCH = 2
# 指定学习率
cfg.SOLVER.BASE_LR = 0.00025
# 最大迭代次数,在达到这一次数后,模型停止优化
cfg.SOLVER.MAX_ITER = 300
# 进行学习率衰减的步数,这里任务较为简单,没有使用学习率衰减
cfg.SOLVER.STEPS = []
# 设定一些模型内部参数,如需要分类的类别数,ROI-head一次处理的proposal的数量
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
以上是配置文件的选取和参数的修改,这在上次的学习中也提到过,在这不做赘述,可以看看上面代码的注释。
接下来就是本节的重点,模型的训练:
from detectron2.engine import DefaultTrainer
# 根据cfg的输入文件夹名称创建一个文件夹
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
# 实例化一个DefaultTrainer对象,并进行训练
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
很显然,负责训练的是一个DefaultTrainer类的对象,我们来看一下它的源码defaults.py
由于代码比较长,我们只挑类中重要的函数来精读:
# 1.先看一下官方的注释,我用我的理解翻译了一下,觉得有问题可以去看源码~
class DefaultTrainer(TrainerBase):
"""
一个能执行默认训练步骤的训练器,他的运行过程如下:
1.用给予的配置文件中的模型,优化器,数据接口来创建一个“SimpleTrainer”类的实例。
2.当调用resume_or_load()函数时,加载最新的权重或者cfg.MODEL.WEIGHT。
3. 根据配置文件注册一些钩子(hook)函数(参考pytorch的hooks)
这个类别用于简化“标准的模型训练流程”,减少重复的代码,只适用于标准的模型训练。它对训练逻辑有很多先验假设,在新的研究任务中可能会失效。
简单来说,这个类别只适用于常规的模型训练,对于新的任务,可能会需要我们根据实际重写方法。
使用实例:
::
trainer = DefaultTrainer(cfg)
trainer.resume_or_load() # load last checkpoint or MODEL.WEIGHTS
trainer.train()
"""
# 2. 初始化函数
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
super().__init__()
logger = logging.getLogger("detectron2")
if not logger.isEnabledFor(logging.INFO): # setup_logger is not called for d2
setup_logger()
cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size())
# 使用build_model(cfg)函数构建模型,与此函数相关的知识见学习笔记第一章
model = self.build_model(cfg)
# 创建优化器
optimizer = self.build_optimizer(cfg, model)
# 创建数据接口
data_loader = self.build_train_loader(cfg)
# 创建分布式训练模型
model = create_ddp_model(model, broadcast_buffers=False)
#创建训练器
self._trainer = (AMPTrainer if cfg.SOLVER.AMP.ENABLED else SimpleTrainer)(
model, data_loader, optimizer
)
# 学习率衰减
self.scheduler = self.build_lr_scheduler(cfg, optimizer)
# 设置检查点
self.checkpointer = DetectionCheckpointer(
# Assume you want to save checkpoints together with logs/statistics
model,
cfg.OUTPUT_DIR,
trainer=weakref.proxy(self),
)
# 设定最小、最大训练轮数
self.start_iter = 0
self.max_iter = cfg.SOLVER.MAX_ITER
self.cfg = cfg
# 注册一些hook函数
self.register_hooks(self.build_hooks())
DefaultTrainer的初始化函数中其实还有很多可以深入研究的地方,比如数据接口和优化器的构建过程,在这里先不做详述,以后有机会再来探究。下面看一下我们需要使用的训练启动函数:
# 执行训练函数
def train(self):
"""
Run training.
Returns:
当存在评估过程时,返回评估结果;否则无返回值。
"""
super().train(self.start_iter, self.max_iter)
if len(self.cfg.TEST.EXPECTED_RESULTS) and comm.is_main_process():
assert hasattr(
self, "_last_eval_results"
), "No evaluation results obtained during training!"
verify_results(self.cfg, self._last_eval_results)
return self._last_eval_results
可以看到,DefaultTrainer实际上是使用其父类TrainerBase的train()函数进行训练,来看看TrainerBase的代码(detectron2/engine/train_loop.py):
class TrainerBase:
"""
基础迭代训练器
对模型、优化器等其他参数没有限定。
Attributes:
iter(int): 目前的迭代轮数
start_iter(int): 开始的训练轮数(按惯例最小为0).
max_iter(int): 结束的训练轮数.
storage(EventStorage): 训练过程中开放的事件储存器.
"""
# 初始化
def __init__(self) -> None:
self._hooks: List[HookBase] = []
self.iter: int = 0
self.start_iter: int = 0
self.max_iter: int
self.storage: EventStorage
_log_api_usage("trainer." + self.__class__.__name__)
def register_hooks(self, hooks: List[Optional[HookBase]]) -> None:
"""
为训练器注册一些hook函数,这些函数以它们被注册的顺序执行
Args:
hooks (list[Optional[HookBase]]): list of hooks
"""
hooks = [h for h in hooks if h is not None]
for h in hooks:
assert isinstance(h, HookBase)
h.trainer = weakref.proxy(self)
self._hooks.extend(hooks)
def train(self, start_iter: int, max_iter: int):
"""
Args:
start_iter, max_iter (int): See docs above
"""
logger = logging.getLogger(__name__)
logger.info("Starting training from iteration {}".format(start_iter))
self.iter = self.start_iter = start_iter
self.max_iter = max_iter
with EventStorage(start_iter) as self.storage:
try:
self.before_train()
for self.iter in range(start_iter, max_iter):
self.before_step()
self.run_step()
self.after_step()
self.iter += 1
except Exception:
logger.exception("Exception during training:")
raise
finally:
self.after_train()
我们重点关注TrainerBase训练部分的代码,其核心部分为:
for self.iter in range(start_iter, max_iter):
self.before_step()
self.run_step()
self.after_step()
可以看到,这是一个简单的循环结构,一次迭代要执行三个函数:before_step(),run_step(),以及after_step():
def before_step(self):
self.storage.iter = self.iter
for h in self._hooks:
h.before_step()
def after_step(self):
for h in self._hooks:
h.after_step()
def run_step(self):
raise NotImplementedError
before_step()和aftert_step()都是对hook函数进行的操作,在此暂且不提。重点在于run_step()函数,可以看到这是一个抽象方法,在使用时需要我们自己重写函数,否则会报错。
既然如此,那么DefaultTrainer中是如何进行训练的呢?我们看到DefaultTrainer中重写的run_step()函数,看一看它用的是怎样的训练方式:
'''
DefaultTrainer中重写了TrainerBase中的抽象方法
显然,它使用self._trainer.run_step()做为自己的run_step()过程
'''
def run_step(self):
self._trainer.iter = self.iter
self._trainer.run_step()
在DefaultTrainer的初始化函数中,对self._trainer做如下的初始化操作:
self._trainer = (AMPTrainer if cfg.SOLVER.AMP.ENABLED else SimpleTrainer)(
model, data_loader, optimizer
)
显然,默认的训练器只可能有两种:AMPTrainer和SimpleTrainer。为了方便初学者学习,暂时不考虑前者,只看看SimpleTrainer的源代码文件(与TrainerBase类在相同的代码文件中),重点关注其run_step()函数:
def run_step(self):
"""
Implement the standard training logic described above.
"""
assert self.model.training, "[SimpleTrainer] model was changed to eval mode!"
start = time.perf_counter()
# 通过迭代器获取数据
data = next(self._data_loader_iter)
data_time = time.perf_counter() - start
# 计算损失函数
loss_dict = self.model(data)
if isinstance(loss_dict, torch.Tensor):
losses = loss_dict
loss_dict = {"total_loss": loss_dict}
else:
losses = sum(loss_dict.values())
# 梯度下降
self.optimizer.zero_grad()
losses.backward()
self._write_metrics(loss_dict, data_time)
self.optimizer.step()
对于使用过pytorch的同学来说,上面的代码应该非常好理解,就是pytorch中最常见的计算损失-梯度重置-反向传播-梯度下降这一套模型训练流程(看到这就突然感觉很熟悉很温馨啊~)
下面我们回到最初的训练代码,对这一小节的内容进行总结:
from detectron2.engine import DefaultTrainer
# 根据cfg的输入文件夹名称创建一个文件夹
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
# 实例化一个DefaultTrainer对象,并进行训练
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
模型的训练步骤为:
- 实例化一个DefaultTrainer对象,并用配置文件(cfg)对其进行初始化.
- 调用trainer.train()进行训练。在默认情况下,训练器将使用SimpleTrainer,即使用pytorch最基础的训练逻辑和训练流程进行训练。
- 训练从start_iter开始,到max_iter结束,这两个参数可以在配置文件中指定,例如:
cfg.SOLVER.MAX_ITER = 300
3.模型评估(Evaluation)
模型训练完之后,我们还要进行模型的推理(Inference)和评估。模型推理部分在第一章的学习笔记中已经做了分析,在此不再赘述,本小节重点讲一讲模型的评估。对于气球的检测与分割任务,我们使用通用的COCO AP来计算模型的得分,代码如下:
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
# 创建Evaluator
evaluator = COCOEvaluator("balloon_val", ("bbox", "segm"), False, output_dir="./output/")
# 构建测试数据集
val_loader = build_detection_test_loader(cfg, "balloon_val")
#打印评估结果
print(inference_on_dataset(trainer.model, val_loader, evaluator))
代码非常简单,只有三行,我们来分别看看对应的源代码。
首先是测试数据的构建(detection2/data/build.py):
其次是COCOEvaluator(detectron2/evaluation/coco_evaluation.py):
def build_detection_test_loader(dataset, *, mapper, sampler=None, num_workers=0):
"""
Args:
dataset (list or torch.utils.data.Dataset):
一个由数据字典组成的列表或者映射式pytorch数据集。可以由函数DatasetCatalog.get(name)或者
get_detection_dataset_dicts获取
mapper (callable):
一个可调用的函数,用于从数据集中提取一个样本、转化为模型可用的形式并返回。
当使用配置文件时,默认的选择为:DatasetMapper(cfg, is_train=False)
sampler (torch.utils.data.sampler.Sampler or None):
产生用于数据集的索引,默认为:class:`InferenceSampler
Returns:
DataLoader: 返回一个torch的DataLoader对象
Examples:
::
data_loader = build_detection_test_loader(
DatasetRegistry.get("my_test"),
mapper=DatasetMapper(...))
# 或者,使用配置文件cfgNode进行实例化:
data_loader = build_detection_test_loader(cfg, "my_test")
"""
if isinstance(dataset, list):
dataset = DatasetFromList(dataset, copy=False)
if mapper is not None:
dataset = MapDataset(dataset, mapper)
if sampler is None:
sampler = InferenceSampler(len(dataset))
# 每个worker在一次测试时只使用一张图片
batch_sampler = torch.utils.data.sampler.BatchSampler(sampler, 1, drop_last=False)
# 创建pytorch形式的dataloader并返回
data_loader = torch.utils.data.DataLoader(
dataset,
num_workers=num_workers,
batch_sampler=batch_sampler,
collate_fn=trivial_batch_collator,
)
return data_loader
可以看到,build_detection_test_loader函数最终返回的是一个常见的Dataloader类的数据加载器。
接下来看一下evaluator(detectron2/evaluation/coco_evaluation.py):
class COCOEvaluator(DatasetEvaluator):
'''
为检测、分割、关键点检测等任务计算COCO 形式的AP值
'''
def __init__(
self,
dataset_name,
tasks=None,
distributed=True,
output_dir=None,
*,
use_fast_impl=True,
kpt_oks_sigmas=(),
):
如上代码所示,初始化一个COCOEvaluator实例需要数据名称、指定任务、是否分布式训练(bool)以及输出的文件夹等参数。至于我们究竟如何使用evaluator,还要看一下inference_on_dataset函数的源代码(detectron2/data/build.py):
def inference_on_dataset(
model, data_loader, evaluator: Union[DatasetEvaluator, List[DatasetEvaluator], None]
):
"""
在指定的数据集、评估器上对于模型性能进行评估
Returns:
此函数返回调用evaluator.evaluate()的返回值
"""
''''''(一些不重要的代码)
results = evaluator.evaluate()
if results = None:
results = {}
return results
上述函数本身没有定义评估的操作,而是执行我们指定的evaluator的evaluate方法。细心的同学可能会发现,之前提到的COCOEvaluator继承自父类DatasetEvaluator,我们来看看父类的evaluate方法,源代码文件同inference_on_dataset()一致:
class DatasetEvaluator:
def evaluate(self):
"""
Evaluate/summarize the performance, after processing all input/output pairs.
Returns:
dict:
我们希望输出的字典有如下形式:
* key: 任务名称 (e.g., bbox)
* value: 一个{评价指标名称:得分}形式的字典, 如: {"AP50": 80}
"""
pass
不出所料,父类中的evaluate也是一个抽象方法,需要我们自己定义要如何对模型的表现进行评估。那么,最终我们需要看的就是在COCOEvaluator类里重写的evaluate()函数:
def evaluate(self, img_ids=None):
"""
Args:
img_ids: a list of image IDs to evaluate on. Default to None for the whole dataset
"""
......(一些不重要的代码)
self._results = OrderedDict()
if "proposals" in predictions[0]:
self._eval_box_proposals(predictions)
if "instances" in predictions[0]:
self._eval_predictions(predictions, img_ids=img_ids)
# Copy so the caller can do whatever with results
return copy.deepcopy(self._results)
上述代码中使用的self._eval_box_proposals()与 self._eval_predictions函数不做赘述,想要探究的同学可以自己去看看代码。实际上这两个函数就是对模型得到的输出进行格式处理,再调用cocoAPI进行模型评估,相信大家对COCOAPI都是熟悉的。
简而言之,以COCO AP的计算为例,模型的评估可以分为三个步骤:
- 使用build_detection_test_loader函数构建测试集。
- 创建一个COCOEvaluator
- 使用inference_on_dataset函数进行评估,并打印结果。
总结
本章主要从源代码入手,学习了如何在detectron2中注册自己的数据集并进行训练、评估。这两章的内容都是基于官方的colab教程,代码是以命令行的形式出现,可能在实践过程中应用价值不大。下一章我将使用detectron2在服务器的Linux系统上中从零开始训练较大规模检测数据集,学习如何在一个项目中使用detectron2,希望能对大家有所帮助。
(我的毕设终于完成啦~ 祝自己毕业快乐~)