- 深度学习
- 后端
- 开发工具
【免费下载链接】tpu
Reference models and tools for Cloud TPUs.
本文以 TensorFlow Cloud TPU 参考模型仓库中的 ViLD 项目文档 为主体,系统讲解 ViLD(Vision and Language Distillation)开放词汇目标检测方法的原理、代码组织、配置参数、数据预处理与推理流程。读完本文,你将掌握如何在当前仓库中复现 ViLD 的 LVIS 评测管线,理解教师-学生蒸馏在检测器中的落地方式,并能够将预训练模型直接迁移到 PASCAL VOC、COCO、Objects365 等任意数据集上做零微调推理。
一、ViLD 要解决的问题:开放词汇目标检测
传统目标检测器只能识别训练数据集中出现过的固定类别集合。开放词汇(open-vocabulary)目标检测的目标是:检测任意文本描述所对应的目标。其根本困难在于训练数据的匮乏——现有检测数据集通常只有数百个类别,且标注新类别成本极高。
ViLD 提出的解决方案是一种视觉-语言知识蒸馏框架,其核心思想是:
- 使用一个预训练的开放词汇图像分类模型(教师模型,本项目采用 OpenAI CLIP)来编码类别文本与候选区域(object proposals)的图像区域;
- 训练一个两阶段检测器(学生模型,本项目基于 Mask R-CNN 架构),使检测框的区域嵌入(region embeddings)与教师模型推断出的文本嵌入、图像嵌入对齐。
从源码结构看,这一思路在仓库中落地为四类模块:蒸馏数据解析(vild_parser.py)、模型结构(vild_model.py)、检测头(vild_head.py)与损失函数(vild_losses.py)。
论文作者在 LVIS 上通过将所有稀有类(rare categories)作为训练时未见过的 novel 类别来构建评测基准,据论文报告 ViLD 以 ResNet-50 骨干取得 16.1 mask APr,甚至超过有监督对照 3.8;模型无需微调即可直接迁移到其他数据集,在 PASCAL VOC、COCO 与 Objects365 上分别取得 72.2 AP50、36.6 AP 与 11.8 AP;在 COCO 上较此前最优方法在 novel AP 与整体 AP 上分别提升 4.8 与 11.4。
二、代码组织与启动入口
ViLD 作为 Cloud TPU detection 框架下的一个项目,代码集中在 models/official/detection/projects/vild:
projects/vild/ ├── README.md # 本文主体对应的项目文档 ├── ViLD_demo.ipynb # Colab 交互演示(两个示例,可换自己的图与类别) ├── requirements.txt # Python 依赖 ├── configs/ │ ├── vild_config.py # 默认参数模板(继承 DETECTION_CFG) │ ├── vild_resnet.yaml # ViLD(vanilla 蒸馏)配置 │ └── vild_ensemble_resnet.yaml # ViLD-ensemble(double_branch 蒸馏)配置 ├── dataloader/ │ ├── tf_example_decoder.py # 解码含蒸馏标签的 tf.Example │ └── vild_parser.py # 图像/标注/蒸馏标签解析器 ├── modeling/ │ ├── vild_head.py # ViLD Fast R-CNN 头(CLIP 分类器 + 特征蒸馏投影) │ ├── vild_losses.py # 分类损失(支持 mask_rare) │ └── vild_model.py # ViLD 模型主体(继承 base_model.BaseModel) ├── ops/ │ └── postprocess_ops.py # 检测后处理(NMS 等) └── preprocessing/ ├── create_lvis_tf_record.py # LVIS → TFRecord 转换脚本 └── dataset_util.pyViLD 通过三处工厂函数挂载到统一框架中:
- 参数模板:configs/factory.py 中
model == 'vild'时返回vild_config.VILD_CFG; - 模型构建:modeling/factory.py 中
params.type == 'vild'时实例化vild_model.ViLDModel; - 数据解析:dataloader/factory.py 中
parser == 'vild_parser'时实例化vild_parser.Parser; - 检测头:modeling/architecture/factory.py 中的
vild_fast_rcnn_head_generator。
启动入口是统一的 models/official/detection/main.py,通过--model="vild"与--mode=eval/train等参数驱动。
三、核心实现:特征蒸馏如何在两阶段检测器内落地
3.1 模型前向流程(vild_model.py)
ViLDModel继承base_model.BaseModel,整体保留 Mask R-CNN 的 ResNet-FPN-RPN 管线,关键改造点(见 vild_model.py):
- 训练时拼接蒸馏 ROI:在 RPN 采样得到
rpn_rois后,若开启特征蒸馏(visual_feature_distill非 None),将教师模型预计算的labels['roi_boxes'](形状[batch, max_distill_rois, 4])沿 axis=1 与 RPN 的 ROI 拼接,随后统一经过multilevel_crop_and_resize提取 ROI 特征,前向结束后再tf.split还原。max_num_rois(默认 300)决定了每张图蒸馏 ROI 的数量; - RoI 特征 → 检测头:
_frcnn_head_fn返回class_outputs / box_outputs / distill_feat_outputs / distill_class_outputs四路输出,其中distill_feat_outputs只在训练期参与蒸馏损失; - 推理期:
_generate_detections_fn使用postprocess_ops.GenericDetectionGenerator做 NMS 等后处理(vild_model.py),ViLD-ensemble 额外传入distill_class_outputs做双分支分类集成。
3.2 蒸馏损失(build_losses)
总损失为(vild_model.py):
model_loss = rpn_score_loss + rpn_box_loss + frcnn_class_loss + frcnn_box_loss + mask_loss + feat_distill_weight * feat_distill_loss其中特征蒸馏损失的计算细节(vild_model.py):
- 训练时将教师预计算的
gt_visual_feat按 L2 范数归一化(normalize_feat_during_training=True); - 用
valid_rois(roi_boxes最大值大于 0 的条目)作为权重掩码,仅对有效 ROI 计算 L1 损失(tf.losses.absolute_difference,SUM_BY_NONZERO_WEIGHTS); - 损失值再乘以
visual_feature_dim(512)以保持量纲稳定; feat_distill_weight为蒸馏损失的缩放系数,配置中 ResNet-50 取 0.5,ResNet-152 取 1.0 或 2.0(见后文 checkpoint 表)。
3.3 两种蒸馏模式:vanilla 与 double_branch
visual_feature_distill支持三种取值(vild_config.py 中注释为None / 'vanilla' / 'double_branch'):
- None:关闭蒸馏,退化为普通分类头;
- vanilla(对应 vild_resnet.yaml):RPN ROI 与蒸馏 ROI 共享同一个 Fast R-CNN 特征提取网络,训练时在投影到 CLIP 空间后按
[num_rois - max_distill_rois, max_distill_rois]切分,前者用于分类,后者作为distill_feat_outputs; - double_branch(对应 vild_ensemble_resnet.yaml):在
ViLDFastrcnnHead中为蒸馏 ROI 建立独立的 conv/fc 分支(distill_conv_{i}、distill_fc{i+6})与独立的投影层(distill-project-to-clip),训练、推理两路各自前向,推理时将蒸馏分支的分类输出distill_class_outputs并入检测结果,实现 ViLD-ensemble 集成效果(vild_head.py 与 vild_head.py)。
3.4 ViLD 检测头与 CLIP 分类器(vild_head.py)
ViLDFastrcnnHead(vild_head.py)的关键参数与行为:
| 参数 | 默认值 | 作用 |
|---|---|---|
clip_dim | 512 | CLIP 文本/视觉嵌入维度 |
classifier_weight_path | 必填 | 预计算的 CLIP 文本嵌入.npy路径,作为分类权重 |
normalize_classifier | True | 是否对分类器权重按列做 L2 归一化后除到 logits 上 |
normalize_visual | True | 是否对投影后的视觉特征做 L2 归一化 |
temperature | 100.0 | 乘到 logits 上的温度系数,放大对比度 |
visual_feature_distill | None | 蒸馏模式(None/vanilla/double_branch) |
max_distill_rois | 300 | 蒸馏 ROI 数量上限 |
class_agnostic_bbox_pred | True | 类别无关的框回归(只预测 4 个坐标而非 num_classes×4) |
实现要点:
- 训练期通过
tf.gfile.GFile(self._classifier_weight_path)加载文本嵌入,断言形状为(clip_dim, num_classes - 1),并以tf.initializers.constant作为class-predict全连接层(无 bias)的核初始化(vild_head.py)。这意味着1204 个类别 = 1 个背景类 + 1203 个 LVIS 类别; - 分类器与背景分支分别做 L2 归一化后拼接(
background_output在前,与tf.concat组合),最后统一乘以temperature(vild_head.py); - 变量作用域为
frcnn_layer_0/fast_rcnn_head,配合配置frozen_variable_prefix: frcnn_layer_0/fast_rcnn_head/class-predict在训练时冻结 CLIP 分类器权重,仅微调视觉投影与检测头。
3.5 分类损失与稀有类掩码(vild_losses.py)
FastrcnnClassLoss(vild_losses.py)在配置frcnn_class_loss.mask_rare=True时,从rare_mask_path加载 LVIS 稀有类二值掩码,将稀有类 logits 置零后计算 softmax 交叉熵。这一设计保证在仅监督基类(frequent/common categories)的训练设定下,稀有类(novel categories)不会产生错误的梯度信号,与论文的“稀有类全部视作 novel”基准保持一致。
四、配置文件深度解析
ViLD 提供两个 YAML 配置,均以 vild_config.py 中的VILD_CFG为默认模板,通过--config_file加载、--params_override覆盖。
4.1 vild_resnet.yaml(ViLD,ResNet-50)
核心字段含义:
- architecture:
backbone: resnet、multilevel_features: fpn,特征金字塔 level 2–6;num_classes: 1204;include_mask: true、mask_target_size: 28(输出 mask 尺寸);visual_feature_distill: vanilla、visual_feature_dim: 512、max_num_rois: 300、feat_distill_weight: 0.5、normalize_feat_during_training: true、filter_distill_boxes_size: 0(0 表示不过滤小于指定尺寸的蒸馏框); - frcnn_head:
num_convs: 4、num_filters: 256、num_fcs: 2、fc_dims: 1024,clip_dim: 512、normalize_classifier: true、normalize_visual: true、temperature: 100.0、class_agnostic_bbox_pred: true; - resnet:
resnet_depth: 50; - frcnn_class_loss:
mask_rare: true(配合推理脚本中的rare_mask_path); - roi_proposal / roi_sampling:RPN 训练期
rpn_pre_nms_top_k: 2000、rpn_post_nms_top_k: 1000、nms_iou_threshold: 0.7;ROI 采样num_samples_per_image: 512、fg_fraction: 0.25、fg_iou_thresh: 0.5、mix_gt_boxes: true; - postprocess:
max_total_size: 300、pre_nms_num_boxes: 1000、nms_iou_threshold: 0.5、score_threshold: 0.0、apply_nms: true; - train:
frozen_variable_prefix: frcnn_layer_0/fast_rcnn_head/class-predict(冻结 CLIP 文本分类器),l2_weight_decay: 4.0e-05,total_steps: 180000,train_batch_size: 256,step 式学习率(init_learning_rate: 0.32、learning_rate_steps: [162000, 171000, 175500]、warmup_steps: 1000),momentum 优化器(0.9); - vild_parser:输入图
output_size: [1024, 1024],RPN 匹配阈值 0.7/0.3,随机水平翻转,尺度增强aug_scale_min: 0.1、aug_scale_max: 2.0,mask_crop_size: 112、max_num_instances: 300,训练时跳过 crowd 标注; - eval:
type: lvis_box_and_mask、eval_samples: 19809、use_json_file: true。
4.2 vild_ensemble_resnet.yaml(ViLD-ensemble)
与 vanilla 配置的差异集中在两点:visual_feature_distill: double_branch(见 3.3 节的双分支实现)与use_tpu: false。README 中 ViLD-ensemble 与 ResNet-152 系列的 checkpoint 均引用该配置,蒸馏权重依据下表使用feat_distill_weight覆盖。
五、环境准备与依赖
运行前需要满足两个前提(详见 README.md):
- 安装 TensorFlow(代码基于
tensorflow.compat.v1编写,需确保导入兼容); - 安装 requirements.txt 中的依赖:
tensorflow pycocotools lvis opencv-python此外,脚本内部存在sys.path.insert(0, 'tpu/models')的路径假设(见 main.py 与 vild_config.py),因此从仓库根目录出发,命令中的配置路径与模块路径应以tpu/models/official/detection/...形式书写,与官方 README 中的脚本保持一致。
六、数据预处理:LVIS v1.0 → TFRecord
6.1 目录结构要求
下载并解压 LVIS v1.0 验证集到DATA_DIR,需按如下结构组织:
DATA_DIR +-- lvis_v1_val.json +-- val2017 | +-- ***.jpg | +-- ...6.2 转换脚本
在models/official/detection/projects/vild/preprocessing/目录下执行(脚本入口为 create_lvis_tf_record.py,max_num_processes可按机器核数调整,DEST_DIR指定 TFRecord 输出目录):
DATA_DIR=[DATA_DIR] DEST_DIR=[DEST_DIR] VAL_JSON="${DATA_DIR}/lvis_v1_val.json" python3 preprocessing/create_lvis_tf_record.py \ --image_dir="${DATA_DIR}" \ --json_path="${VAL_JSON}" \ --dest_dir="${DEST_DIR}" \ --include_mask=True \ --split='val' \ --num_parts=100 \ --max_num_processes=100脚本支持的 flags 包括--include_mask(是否包含 PNG 编码的实例分割掩码)、--split(取值train/val)、--num_parts(生成的 TFRecord 分片数)、--debug等。脚本会把 LVIS 的[x, y, width, height]绝对坐标转换为 TensorFlow Object Detection API 期望的[ymin, xmin, ymax, xmax]归一化格式,并利用 pycocotools 的 mask 工具解码 RLE 分割标注。
6.3 训练数据说明
从 vild_parser.py 与 tf_example_decoder.py 可以看出,ViLD 训练所用的 TFRecord 除常规图像、框、类别、掩码标注外,还须包含教师模型预计算的蒸馏标签:roi_boxes(蒸馏候选框)与groundtruth_visual_features(CLIP 图像区域嵌入)。解析器在训练期会将这些标签随 RPN 标签一起打包(labels['roi_boxes']、labels['gt_visual_feat'],clip 到max_num_rois长度),并可选地按filter_distill_boxes_size过滤过小的蒸馏框。这类带蒸馏标签的训练数据需要依据论文的离线教师推理流程预先产生,官方文档未在仓库内给出其生成脚本。
七、预训练 Checkpoint 与评测指标
官方在 LVIS(稀有类全部作为 novel)上发布的预训练模型与指标如下表(数据来自 README.md):
| Method | Backbone | Distillation weight | APr | APc | APf | AP | config | ckpt |
|---|---|---|---|---|---|---|---|---|
| ViLD | ResNet-50 | 0.5 | 16.6 | 19.8 | 28.2 | 22.5 | vild_resnet.yaml | resnet50_vild.tar.gz |
| ViLD-ensemble | ResNet-50 | 0.5 | 18 | 24.7 | 30.6 | 25.9 | vild_resnet.yaml | resnet50_vild_ensemble.tar.gz |
| ViLD | ResNet-152 | 1.0 | 19.6 | 21.6 | 28.5 | 24.0 | vild_ensemble_resnet.yaml | resnet152_vild.tar.gz |
| ViLD-ensemble | ResNet-152 | 2.0 | 19.2 | 24.8 | 30.8 | 26.2 | vild_ensemble_resnet.yaml | resnet152_vild_ensemble.tar.gz |
其中 APr/APc/APf 分别对应 LVIS 的稀有类、常见类与频繁类 AP;蒸馏权重(feat_distill_weight)越大,模型越偏向蒸馏信号。官方 checkpoint 均托管于云存储,下载地址见原文档表格。另外据论文报告,使用 ResNet-50 骨干时 ViLD 取得 16.1 mask APr,较有监督对照提升 3.8;跨数据集零微调迁移结果为 PASCAL VOC 72.2 AP50、COCO 36.6 AP、Objects365 11.8 AP。
八、推理(Evaluation)完整流程
8.1 准备推理所需的三个输入
- 下载 CLIP 文本嵌入 clip_synonym_prompt.npy 与稀有类二值掩码 lvis_rare_masks.npy,放入
[WEIGHTS_DIR]; - 下载并解压上表任一预训练模型到
[MODEL_DIR]; - 将
[RESNET_DEPTH]、[MODEL_DIR]、[DATA_DIR]、[DEST_DIR]、[WEIGHTS_DIR]、[CONFIG_FILE]替换为实际值。
8.2 运行脚本
BATCH_SIZE=1 RESNET_DEPTH=[RESNET_DEPTH] MODEL_DIR=[MODEL_DIR] EVAL_FILE_PATTERN="[DEST_DIR]/val*" VAL_JSON_FILE="[DATA_DIR]/lvis_v1_val.json" RARE_MASK_PATH="[WEIGHTS_DIR]/lvis_rare_masks.npy" CLASSIFIER_WEIGHT_PATH="[WEIGHTS_DIR]/clip_synonym_prompt.npy" CONFIG_FILE="tpu/models/official/detection/projects/vild/configs/[CONFIG_FILE]" python3 tpu/models/official/detection/main.py \ --model="vild" \ --model_dir="${MODEL_DIR?}" \ --mode=eval \ --use_tpu=False \ --config_file="${CONFIG_FILE?}" \ --params_override="{ resnet: {resnet_depth: ${RESNET_DEPTH?}}, predict: {predict_batch_size: ${BATCH_SIZE?}}, eval: {eval_batch_size: ${BATCH_SIZE?}, val_json_file: ${VAL_JSON_FILE?}, eval_file_pattern: ${EVAL_FILE_PATTERN?} }, frcnn_head: {classifier_weight_path: ${CLASSIFIER_WEIGHT_PATH?}}, postprocess: {rare_mask_path: ${RARE_MASK_PATH?}}}"8.3 参数含义与底层行为
--model="vild"由 main.py 决定加载 ViLD 参数模板与模型;--mode=eval触发评估;--use_tpu=False时 main.py 会强制关闭 bfloat16 与同步 BN,因此纯 GPU/CPU 即可运行推理;--config_file指向vild_resnet.yaml或vild_ensemble_resnet.yaml,两者通过--params_override覆盖骨干深度、批大小、数据路径与两个.npy权重路径;classifier_weight_path最终进入ViLDFastrcnnHead,作为class-predict层核初始化并在归一化后参与打分;rare_mask_path同时服务于frcnn_class_loss(训练期掩码)与后处理阶段的稀有类约束;- 评估类型为
lvis_box_and_mask,输出 mask 与 box 双模态指标;predict_batch_size与eval_batch_size保持一致即可(示例取 1,以降低显存占用)。
更完整的训练/评估流程可参考检测模块通用指南 GETTING_STARTED.md。
九、Colab 交互演示
仓库提供两种演示途径(见 README.md):
- 在线 Colab:原文档给出的
ViLD_demo.ipynb托管版本; - 本地 Notebook:ViLD_demo.ipynb。
演示内置两个示例,同时支持上传自己的图片并指定想检测的类别文本,直观展示开放词汇检测能力。
十、引用与说明
本项目致谢部分明确说明:教师模型使用 OpenAI 的 CLIP 模型;整体代码构建于 Cloud TPU detection 框架之上。若在研究中引用 ViLD,可按原文档给出的 BibTeX 条目引用(Gu, Xiuye; Lin, Tsung-Yi; Kuo, Weicheng; Cui, Yin. "Open-Vocabulary Detection via Vision and Language Knowledge Distillation." arXiv:2104.13921, 2021.)。仓库许可证与 LICENSE 一致,使用前请留意。
结语
ViLD 通过将 CLIP 的开放词汇能力蒸馏进两阶段检测器,为开放词汇目标检测提供了一个可复现、可迁移的完整实现。本文沿着“原理 → 源码 → 配置 → 数据 → 推理”的路径,覆盖了 vild/README.md 的全部实操内容,并深入到 vild_model.py、vild_head.py、vild_config.py 等实现细节,读者可按文中命令在本地或 Cloud TPU 上复现 LVIS 评测,或将预训练模型直接用于自定义数据集的开放词汇检测。
- 深度学习
- 后端
- 开发工具
【免费下载链接】tpu
Reference models and tools for Cloud TPUs.
相关推荐
F-VLM 实战指南:基于冻结视觉语言模型的开词汇目标检测(JAX/FLAX 实现)
F VLM 实战指南:基于冻结视觉语言模型的开词汇目标检测(JAX/FLAX 实现) F VLM(Frozen Vision Language Model)是
人工智能深度学习NLP计算机视觉强化学习MMDetection 中的 GLIP 实现指南:视觉-语言联合预训练的开放词汇目标检测
MMDetection 中的 GLIP 实现指南:视觉 语言联合预训练的开放词汇目标检测 GLIP(Grounded Language Image Pre tr
人工智能计算机视觉深度学习模型评测Sentence Transformers 稀疏编码器知识蒸馏实战:基于 SparseMarginMSELoss 训练 SPLADE 检索模型
Sentence Transformers 稀疏编码器知识蒸馏实战:基于 SparseMarginMSELoss 训练 SPLADE 检索模型 知识蒸馏是训练最
人工智能NLPEmbedding微调机器学习
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考