☰
Detectron2 数据集层设计解析:最小数据结构与惰性加载如何支撑高效的视觉训练
2026/10/4 1:46:24 网站建设 项目流程
  • 人工智能
  • 计算机视觉
  • 深度学习
  • 预训练
  • 微调

【免费下载链接】Vim

[ICML 2024] Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

项目地址:https://gitcode.com/gh_mirrors/vim2/Vim
点击查看免费下载

导读

本文围绕det/detectron2/data/datasets/目录下数据集模块的核心设计哲学展开:数据集实现只需提供"最小数据结构",无需把数据加载成模型可直接消费的最终格式;对于图像数据集,仅需提供文件名与标签,不要读取图像内容,具体如何读取由下游(训练/推理流水线)决定。读懂这套设计,你就能理解 Detectron2(以及本仓库det/下基于它的 Vision Mamba 检测训练流程)的数据加载性能关键,并掌握如何用同样的模式注册自定义数据集。

一、核心设计原则:最小数据结构 + 惰性加载

det/detectron2/data/datasets/README.md开宗明义地定义了本目录的职责边界,原文要点如下:

The dataset implemented here do not need to load the data into the final format. It should provide the minimal data structure needed to use the dataset, so it can be very efficient. For example, for an image dataset, just provide the file names and labels, but don't read the images. Let the downstream decide how to read.

翻译并拆解成三条可操作的原则:

  1. 不做最终格式的转换:数据集模块的产物不是"已经变成 Tensor / Instances 的样本",而是一份轻量的、Python 字典构成的元数据清单;
  2. 只提供最小数据结构:对图像数据集而言,最小结构就是file_name(图像路径)+ 标注字段(框、类别、分割等);
  3. 惰性读取图像:加载注解时绝不cv2.imread/PIL.Image.open读取像素,把"何时真正读图"完全交给下游的DatasetMapper与 DataLoader。

这一原则的直接收益在 coco.py 的函数注释中写得很明确:"This function does not read the image files. The results do not have the 'image' field."——一个十几万张的 COCO 训练集,注解加载阶段只解析 JSON 与路径字符串,内存开销极小,且所有 worker 可以共享同一份list[dict]。

二、全局注册表:DatasetCatalog 与 MetadataCatalog

"最小数据结构"得以流转的核心,是 catalog.py 中定义的两个全局注册表:

  • DatasetCatalog:字符串(数据集名,如coco_2014_train)→ 一个"无参数、返回list[dict]"的函数。register(name, func)注册、get(name)调用函数拿到注解列表、list()列出所有已注册名称。它的设计目的(见 catalog.py)是"通过配置中的字符串即可方便地切换数据集";
  • MetadataCatalog:字符串 →Metadata单例对象,用于存储常量型共享知识(如 COCO 的类别名列表、可视化颜色、评估器类型)。它提供set(**kwargs)批量写入与get(key, default)安全读取,且同名 metadata 一旦创建便常驻(catalog.py)。

两个注册表是后续所有数据集加载器、DatasetMapper、build_detection_train_loader(见 build.py)共用的全局枢纽:数据集加载器只负责往DatasetCatalog注册"产最小字典"的函数,训练流水线则通过DatasetCatalog.get(name)惰性触发。

三、预定义数据集:builtin.py 的注册清单

buitlin.py(实际文件名为builtin.py)以硬编码方式注册了常见数据集的路径与元数据,其头注释强调两个意图:

  1. 加载数据集时能做一致性检查;
  2. 用户无需下载注解文件(JSON)就能直接在标准数据集上使用预训练模型、跑 demo(因为类别名与颜色已被硬编码在元数据中)。

该文件通过_PREDEFINED_SPLITS_COCO(COCO 检测/关键点)、_PREDEFINED_SPLITS_COCO_PANOPTIC(全景分割)、_PREDEFINED_SPLITS_LVIS、_RAW_CITYSCAPES_SPLITS、PASCAL VOC 的SPLITS列表以及 ADE20K 的 train/val 划分,统一在文件末尾的入口代码中注册:

if __name__.endswith(".builtin"): _root = os.path.expanduser(os.getenv("DETECTRON2_DATASETS", "datasets")) register_all_coco(_root) register_all_lvis(_root) register_all_cityscapes(_root) register_all_cityscapes_panoptic(_root) register_all_pascal_voc(_root) register_all_ade20k(_root)

值得注意的注册细节:

  • 默认数据根目录约定:所有预定义数据集默认假设存放在./datasets/下,可用环境变量DETECTRON2_DATASETS覆盖(builtin.py);
  • 支持远程路径:register_all_coco中os.path.join(root, json_file) if "://" not in json_file else json_file(builtin.py)说明注解 JSON 可以是带://的远程 URI,由PathManager统一处理;
  • 全景分割的双注册:同一前缀同时注册register_coco_panoptic_separated(供 Panoptic FPN 使用)与register_coco_panoptic(供 Panoptic-DeepLab 使用)两种格式(builtin.py);
  • COCO 类别 id 约定:COCO 的 category id 存在人为删除导致不连续的问题,注册时会生成thing_dataset_id_to_contiguous_id映射到连续 id(详见下文load_coco_json)。

元数据本身存放在 builtin_meta.py 中,例如COCO_CATEGORIES每项包含color(可视化颜色)、isthing、id、name。其注释同样强调:自定义数据集无需在代码中硬编码元数据——COCO 格式数据集调用load_coco_json时元数据会自动生成,硬编码只是为了"没有下载 JSON 的用户也能直接可视化模型"。

四、各加载器如何产出"最小数据结构"

4.1 COCO:load_coco_json(实例检测/分割/关键点)

核心函数位于 coco.py,签名与行为:

def load_coco_json(json_file, image_root, dataset_name=None, extra_annotation_keys=None):
  • 内部用pycocotools.COCO解析 JSON,只保留每个样本的字段:file_name、height、width、image_id、annotations(每项含iscrowd、bbox、keypoints、category_id、segmentation,以及extra_annotation_keys指定的额外键);
  • 当传入dataset_name时,会顺带把元数据写入MetadataCatalog:根据 JSON 的categories字段生成thing_classes,并将category_id映射到连续区间[0, #categories)生成thing_dataset_id_to_contiguous_id(coco.py)。若类别 id 本就连续则免去映射,否则打印警告;
  • 对分割做基本清洗:过滤少于 3 个点的无效多边形、把list形式的 RLE 转成压缩 RLE、为关键点坐标加 0.5(从像素索引转为浮点坐标,coco.py);
  • 注册入口register_coco_instances(name, metadata, json_file, image_root)(coco.py)示范了标准的"两步注册法":
# 1. 注册一个返回最小字典列表的函数(惰性:调用时才解析) DatasetCatalog.register(name, lambda: load_coco_json(json_file, image_root, name)) # 2. 附加元数据(评估器类型、路径等) MetadataCatalog.get(name).set( json_file=json_file, image_root=image_root, evaluator_type="coco", **metadata )

4.2 PASCAL VOC:load_voc_instances

pascal_voc.py 读取ImageSets/Main/{split}.txt得到文件 id 列表,逐个解析Annotations/*.xml,输出同样的最小字典:file_name、image_id、height、width、annotations(category_id用CLASS_NAMES.index(cls)得到连续 id,bbox_mode=BoxMode.XYXY_ABS)。注意它把注解的 1 基像素坐标统一减 1 转为 0 基坐标(pascal_voc.py),并保留了 "difficult" 样本参与训练——这是作者基于实验的选择(pascal_voc.py 中注释掉的过滤代码)。

4.3 语义分割:load_sem_seg

coco.py 中的load_sem_seg(gt_root, image_root, gt_ext="png", image_ext="jpg")面向语义分割:扫描image_root与gt_root下按扩展名匹配的文件,依据相对路径(不含扩展名)匹配图像与标注,每个样本只输出file_name与sem_seg_file_name两个字段——同样不读像素。它还处理了图像数与标注数不一致的情况:取两者 basename 的交集并排序,保证多 worker 下顺序一致(coco.py)。

4.4 Cityscapes:多进程预处理

cityscapes.py 的load_cityscapes_instances是"预处理较重"的案例:注解以gtFine目录下的 png/json 形式存放,解析需要把 instance mask 转成多边形。源码用mp.Pool多进程并行处理,并给出性能提示:"所有 worker 会重复这份工作,在 8 GPU 服务器上可能耗时 10 分钟"(cityscapes.py)。这说明"最小数据结构"的边界取决于数据源成本:Cityscapes 必须解析掩码才能得到实例标注,但读图这一步依然被推迟到下游。

五、下游如何真正读取图像:DatasetMapper 与 DataLoader

"最小数据结构"在训练/推理阶段被 dataset_mapper.py 中的DatasetMapper消费。它的 docstring 明确列出了三步职责(dataset_mapper.py):

  1. 从"file_name"读取图像;
  2. 对图像与标注应用裁剪/几何变换;
  3. 把数据与标注组织成 Tensor 与Instances。

__call__的入口代码(dataset_mapper.py)正是"惰性"兑现的时刻:

dataset_dict = copy.deepcopy(dataset_dict) # 会被原地修改,先深拷贝 image = utils.read_image(dataset_dict["file_name"], format=self.image_format) utils.check_image_size(dataset_dict, image) if "sem_seg_file_name" in dataset_dict: sem_seg_gt = utils.read_image(dataset_dict.pop("sem_seg_file_name"), "L").squeeze(2) ... dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1)))

也就是说:数据集加载器只给路径,DatasetMapper才真正把像素读进内存并转成 Tensor。这样注解解析只在DatasetCatalog.get()时执行一次,而图像读取被推迟到 DataLoader 采样阶段,且可被 PyTorch 多进程 worker 并行分摊。

数据流全貌(依据 build.py 与 common.py 的调用关系):

配置 DATASETS.TRAIN/TEST(字符串) → DatasetCatalog.get(name) # 惰性调用注册函数,产出 list[dict](最小结构) → DatasetFromList / MapDataset # 包装为 torch Dataset,map 函数即 DatasetMapper → DatasetMapper.__call__ # 此刻才 read_image + 增强 + 转 Tensor/Instances → 采样器(TrainingSampler 等)+ DataLoader → 模型

common.py 中的MapDataset值得额外说明:它把DatasetMapper作为map_func逐样本应用,并且map_func返回None时可跳过坏样本(map-style 下随机换索引重试,common.py);PicklableWrapper则保证 lambda 形式的映射函数也能被多进程 worker pickle。

六、在 Vim(Vision Mamba)检测流程中的实际应用

本仓库的det/是基于 detectron2 的检测工程(对应 ICML 2024 论文《Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model》),训练脚本 det/detectron2/engine/defaults.py 与 det/detectron2/data/build.py 组合使用上面这套机制:

  • 配置文件(如 det/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml)通过DATASETS.TRAIN/DATASETS.TEST字符串引用注册表中的数据集名,切换数据集只需改字符串——这正是DatasetCatalog被设计出来的目的;
  • Vim 骨干网络(det/detectron2/modeling/backbone/vim.py)与 COCO 数据集之间不产生任何直接耦合:数据层只关心"最小字典 → 图像 Tensor",骨干网络只消费 Tensor,二者由build_detection_train_loader与模型 forward 衔接。

因此,即使要在这个工程里引入新的检测数据集(例如自定义实例分割数据),也不需要改动任何模型代码:按第四节register_coco_instances的"两步注册法"(或直接调用load_coco_json返回的 dict 列表)注册新名称,再把配置里的DATASETS.TRAIN/TEST指向它即可。

七、自定义数据集注册的推荐姿势

综合上述源码,注册一个符合本设计哲学的自定义数据集,应遵循:

  1. 写一个返回list[dict]的解析函数:每个 dict 至少包含file_name(或sem_seg_file_name)、height、width、image_id,实例级任务还需annotations列表(每项含bbox+bbox_mode+category_id,可选segmentation/keypoints/iscrowd);
  2. 解析阶段绝不读图像:只拼接路径、解析标注文件(JSON/XML/PNG 掩码等);
  3. 用DatasetCatalog.register+MetadataCatalog.get(name).set(...)注册,并建议像register_coco_instances那样把函数用lambda包一层,让真正的解析延后到首次get时执行;
  4. 元数据尽量自动化:若注解是 COCO 格式,直接复用load_coco_json(json_file, image_root, dataset_name)即可自动获得thing_classes与 id 连续化映射(coco.py);否则参考 builtin_meta.py 的_get_builtin_metadata补充thing_classes、evaluator_type等字段;
  5. 注册时机:参照builtin.py的入口代码,在__main__.endswith(".builtin")条件下(即作为包被导入时)批量执行注册,并把数据根目录做成可配置(如DETECTRON2_DATASETS环境变量),避免硬编码绝对路径。

总结

det/detectron2/data/datasets/README.md用三句话定义了一个影响深远的工程约定:数据集实现只负责产出最小数据结构,不负责把数据加载成最终格式;图像数据集只给文件名与标签,不读图像;读取方式由下游决定。本文以 catalog.py、builtin.py、coco.py、pascal_voc.py、dataset_mapper.py 与 build.py 等源码为证据,还原了这一约定从"注册表 → 惰性加载器 → 下游真正读图"的完整链路。理解这套分层,无论是排查数据加载瓶颈、接入新数据集,还是阅读 Vim 检测工程的训练流程,都会事半功倍。

  • 人工智能
  • 计算机视觉
  • 深度学习
  • 预训练
  • 微调

【免费下载链接】Vim

[ICML 2024] Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

项目地址:https://gitcode.com/gh_mirrors/vim2/Vim
点击查看免费下载

相关推荐

上一篇:Redwood 服务器部署入门:SSH 连接远程服务器的完整实战指南
下一篇:Eclipse EDC连接器配置终极指南:从新手到专家的完整解决方案

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询