sd-scripts 数据集标注实战：用 WD14Tagger 自动为训练图像生成 Danbooru 标签-平芜编程栈

深度学习
计算机视觉
媒体生成
模型训练
微调

【免费下载链接】sd-scripts

项目地址：https://gitcode.com/gh_mirrors/sd/sd-scripts

本篇技术指南围绕 sd-scripts 仓库中的 WD14Tagger 标注文档展开，系统讲解如何用finetune/tag_images_by_wd14_tagger.py脚本对 LoRA、DreamBooth、Textual Inversion 等训练任务的图像数据集进行自动打标。读完本文，你将掌握 ONNX 推理环境的搭建、模型自动下载机制、全部命令行参数（含默认值与源码实现细节），并能按 Animagine XL 3.1 等主流打标规范一键生成可直接用于训练的高质量标签文件。

一、WD14Tagger 是什么，为什么需要它

WD14Tagger 是基于 Waifu Diffusion 1.4 数据集训练的 Danbooru 标签自动识别模型（SmilingWolf 模型系列），能够从图像中自动预测出成千上万个 Danbooru 风格的标签（tag），包括角色、画风、质量、评分等类别。在 sd-scripts 的 LoRA / DreamBooth 训练流程中，高质量的文本标注（caption / tag）是决定训练效果的关键因素之一，WD14Tagger 就是官方推荐的批量打标工具。

脚本的实现位于 finetune/tag_images_by_wd14_tagger.py，其核心思路为：

从 Hugging Face 自动下载模型权重与标签映射表（selected_tags.csv或tag_mapping.json）；
对目录下的每张图像做统一预处理后送入模型推理；
将置信度高于阈值的标签按可配置的规则排序、过滤、替换，最终写入与图像同名的.txt文件（或 JSON / JSONL 元数据）。

二、环境准备：安装 ONNX 运行时

官方推荐使用ONNX进行推理（相比 TensorFlow 部署更轻量、性能更好）。安装命令如下：

pip install onnx onnxruntime-gpu

说明：

onnx提供模型加载与解析能力，onnxruntime-gpu是 GPU 加速的推理运行时（脚本源码在 finetune/tag_images_by_wd14_tagger.py 中通过import onnxruntime as ort使用）；
如果不想用 ONNX，脚本也支持 TensorFlow 推理路径（未指定--onnx时启用），此时需要另行安装 TensorFlow，并加载saved_model.pb/keras_metadata.pb等文件（见源码中的FILES常量）；
模型权重不需要手动下载，首次运行脚本时会从 Hugging Face 自动拉取。

三、基本用法：一条命令完成批量打标

脚本的通用调用形式为：

python finetune/tag_images_by_wd14_tagger.py --onnx --repo_id <模型repo id> --batch_size <批大小> <训练数据文件夹>

官方示例：使用SmilingWolf/wd-swinv2-tagger-v3模型、批大小设为 4，训练数据放在上级目录的train_data文件夹：

python finetune/tag_images_by_wd14_tagger.py --onnx --repo_id SmilingWolf/wd-swinv2-tagger-v3 --batch_size 4 ..\train_data

运行要点：

模型下载：首次运行时，模型文件会自动下载到当前目录下的wd14_tagger_model文件夹中（可通过--model_dir修改保存位置）。根据源码，模型仓库为 SmilingWolf 标准结构时下载model.onnx（ONNX 模式）或saved_model.pb、keras_metadata.pb、selected_tags.csv（TensorFlow 模式），其中 TensorFlow 模式的变量文件存放在variables子目录（见 finetune/tag_images_by_wd14_tagger.py）；
输出位置：标签文件与训练图像放在同一目录下，文件名为图像名 +.txt扩展名（可用--caption_extension修改）；
图像收集：脚本通过 library/dataset.py 的glob_images_pathlib查找图片，支持--recursive递归扫描子文件夹；
推理前处理：源码中的preprocess_image（finetune/tag_images_by_wd14_tagger.py）会将带透明通道的图像合成到白色背景、RGB 转 BGR、居中补边成正方形，再统一缩放到448×448（常量IMAGE_SIZE）。

3.1 面向 Animagine XL 3.1 的打标示例

原文档提供了一个针对Animagine XL 3.1数据规范的完整示例（实际使用时输入为单行命令）：

python finetune/tag_images_by_wd14_tagger.py --onnx --repo_id SmilingWolf/wd-swinv2-tagger-v3 --batch_size 4 --remove_underscore --undesired_tags "PUT,YOUR,UNDESIRED,TAGS" --recursive --use_rating_tags_as_last_tag --character_tags_first --character_tag_expand --always_first_tags "1girl,1boy" ..\train_data

这段命令综合运用了多种标签编辑能力：

--remove_underscore：把标签中的下划线替换为空格（如black_hair→black hair），符合自然语言描述习惯；
--undesired_tags：剔除指定标签（按需替换PUT,YOUR,UNDESIRED,TAGS）；
--recursive：递归处理子目录；
--use_rating_tags_as_last_tag：把评分标签（rating）放到标签序列末尾；
--character_tags_first：角色标签置于最前；
--character_tag_expand：把chara_name_(series)形式的角色标签展开为chara_name, series两个标签；
--always_first_tags "1girl,1boy"：只要图像命中这些标签就强制放到最前面。

3.2 可用的模型仓库 ID

官方支持 SmilingWolf 的 V2、V3 系列模型，按SmilingWolf/wd-vit-tagger-v3的格式指定即可。常用的 V3 模型包括：

SmilingWolf/wd-vit-tagger-v3
SmilingWolf/wd-swinv2-tagger-v3
SmilingWolf/wd-convnext-tagger-v3

默认值：省略--repo_id时使用SmilingWolf/wd-v1-4-convnext-tagger-v2（见源码常量DEFAULT_WD14_TAGGER_REPO，finetune/tag_images_by_wd14_tagger.py）。

此外，从源码结构看（finetune/tag_images_by_wd14_tagger.py），--repo_id还支持namespace/repo_name/subdir这种带子目录的形式，此时脚本会改用model_optimized.onnx+tag_mapping.json的非默认格式加载模型。

四、全部命令行选项详解

所有选项均可通过python tag_images_by_wd14_tagger.py --help查看。以下按类别整理，参数默认值与源码中setup_parser()（finetune/tag_images_by_wd14_tagger.py）保持一致。

4.1 一般选项

选项	默认值	说明
`--onnx`	关闭	使用 ONNX 推理；不指定则走 TensorFlow 路径（需自行安装 TensorFlow）
`--batch_size`	`1`	一次处理的图像数量，按 VRAM 大小调整
`--caption_extension`	`.txt`	输出标签文件的扩展名
`--max_data_loader_n_workers`	不启用	指定 ≥1 时启用 DataLoader 多进程预加载图像，加快读取速度；未指定则同步加载
`--thresh`	`0.35`	输出标签的置信度阈值；调低会得到更多标签，但精度下降
`--general_threshold`	同`--thresh`	一般（General）类标签的独立阈值
`--character_threshold`	同`--thresh`	角色（Character）类标签的独立阈值；设 >1 可禁用角色标签
`--recursive`	关闭	递归处理指定文件夹内的子文件夹
`--append_tags`	关闭	向已存在的标签文件追加标签而非覆盖
`--frequency_tags`	关闭	处理结束后打印各标签的出现频次统计
`--debug`	关闭	输出调试信息（每张图的 rating / character / general / other 标签明细）

4.2 模型下载相关

选项	默认值	说明
`--model_dir`	`wd14_tagger_model`	模型文件的保存目录
`--force_download`	关闭	强制重新下载模型文件

关于批大小，源码还有一个自动校正逻辑：加载 ONNX 模型后会读取其输入张量的 batch 维度，若与--batch_size不一致（且模型为固定批大小），会以模型批大小为准并给出警告（finetune/tag_images_by_wd14_tagger.py）。

推理加速器方面，ONNX 路径会自动探测可用 Provider，优先级为：OpenVINO（GPU + FP32，源码注释说明 fp16 会产生乱码输出）→ CUDA → ROCm → CPU（finetune/tag_images_by_wd14_tagger.py）。

4.3 标签编辑相关

选项	说明
`--remove_underscore`	输出标签中的下划线替换为空格
`--undesired_tags`	从输出中剔除的标签，逗号分隔，如`black eyes,black hair`
`--use_rating_tags`	把评分标签（rating）放在标签序列开头
`--use_rating_tags_as_last_tag`	把评分标签追加到标签序列末尾
`--character_tags_first`	角色标签始终排在一般标签之前
`--character_tag_expand`	展开角色标签的系列名，如`chara_name_(series)`→`chara_name, series`
`--always_first_tags`	指定命中后强制置顶的标签，逗号分隔，如`1girl,1boy`
`--caption_separator`	标签间的分隔字符串，默认`,`（逗号+空格）
`--tag_replacement`	标签替换规则，格式为`源1,目标1;源2,目标2;...`

--tag_replacement的典型应用场景（来自原文档示例）：

想单独学习某套服装：aira tsubase,aira tsubase (uniform)
系列名没有体现在标签中时补全：aira tsubase,aira tsubase\, heir of shadows（其中\,表示标签名内部真实的逗号）

从源码实现看（finetune/tag_images_by_wd14_tagger.py），tag_replacement的解析过程是：先用\转义保护标签名内部的,和;，再按;切分替换对、按,切分源与目标，最后替换命中的标签。若格式错误（非源,目标两段）会直接断言报错。

4.4 标签编辑选项的联动规则（重要）

原文档明确了几条容易踩坑的联动规则：

tag_replacement在character_tag_expand之后应用：即先完成角色标签展开，再进行标签替换；
指定--remove_underscore时：undesired_tags、always_first_tags、tag_replacement中的标签要写成不含下划线的形式（因为输出端已经先做了下划线替换）；
指定--caption_separator时：undesired_tags和always_first_tags要用caption_separator作为分隔符传入，而tag_replacement必须始终用,分隔（源码中undesired_tags、always_first_tags正是按stripped_caption_separator切分的，见 finetune/tag_images_by_wd14_tagger.py）。

4.5 原文档之外的补充选项（源码确认）

python ... --help输出中还包含几个原文档未展开、但实际可用的选项：

--output_path：不写 txt 文件，而是把所有结果汇总输出为 JSON 元数据（若路径以.jsonl结尾则输出 JSONL 格式），每条记录包含image_path、caption（tags 文本）与image_size，便于后续直接供训练管线消费（finetune/tag_images_by_wd14_tagger.py）；
--use_quality_tags/--use_quality_tags_as_last_tag：把 Quality 类标签放在开头 / 末尾（针对非默认格式模型的类别体系）；
--meta_threshold、--model_threshold、--copyright_threshold、--artist_threshold：对 Meta、Model、Copyright、Artist 类别分别设置独立阈值，缺省时均回落到--thresh，设 >1 可禁用对应类别（finetune/tag_images_by_wd14_tagger.py）；
--caption_extention：历史拼写错误的兼容选项，若设置会自动覆盖--caption_extension。

五、源码级原理：标签分类与排序逻辑

理解标签的组织方式有助于你正确使用阈值与排序选项。源码区分了两套标签体系：

SmilingWolf 默认格式（selected_tags.csv）：CSV 每行包含tag_id, name, category, count，其中 category 为9表示评分（rating）、0表示一般（general）、4表示角色（character）（finetune/tag_images_by_wd14_tagger.py）。推理时前 4 个输出对应评分标签，取 argmax 选出置信度最高的评分；其余输出按类别阈值过滤后加入标签列表。

非默认格式（tag_mapping.json）：类别扩展为Rating / General / Character / Copyright / Meta / Model / Quality / Artist八类，输出按概率降序排序，并支持 Quality、Rating 置顶 / 置尾（finetune/tag_images_by_wd14_tagger.py）。

几个值得注意的实现细节：

--remove_underscore只对长度大于 3的标签执行下划线替换，短标签（如1girl）保持原样（finetune/tag_images_by_wd14_tagger.py）；
--append_tags会先读取已有标签文件、去掉重复项，再把新标签追加在末尾（finetune/tag_images_by_wd14_tagger.py）；
--frequency_tags在结束时按出现次数降序打印全部标签统计，可用于检查数据集的标签分布偏差；
--debug会逐图打印 rating / character / general / other 四类标签文本，便于排查打标异常。

六、与训练流程的衔接：标签文件 → 训练元数据

WD14Tagger 生成的.txt标签文件可以直接被 sd-scripts 的多个训练脚本读取（例如--dataset_class对应的文件夹式数据组织）。如果需要把标签合并进 JSON 元数据（metadata）文件，仓库还提供了配套脚本 finetune/merge_dd_tags_to_metadata.py：

python finetune/merge_dd_tags_to_metadata.py <训练数据文件夹> <输出metadata.json> [--in_json 输入] [--full_path] [--recursive]

该脚本会把每个图像对应的标签文本写入 metadata 的tags字段（finetune/merge_dd_tags_to_metadata.py），从而衔接prepare_buckets_latents.py以及 LoRA / DreamBooth 训练的 metadata 模式。关于 metadata 的具体格式规范，可参考 dataset_metadata 说明与 train_README-zh.md。

七、实战建议与注意事项

优先使用 ONNX + GPU 运行时：在 CUDA 环境下脚本会自动选中CUDAExecutionProvider；没有 GPU 时回退到 CPU 推理（此时建议调小--batch_size）；
按显存调节批大小：--batch_size默认 1，VRAM 充足时增大到 4～16 可显著提升吞吐；若与 ONNX 模型固定批大小不一致，脚本会自动采用模型批大小；
善用分类阈值：角色标签常常过拟合误报，可用--character_threshold（设 >1 直接禁用）单独调控；同样支持对 Copyright / Artist / Meta / Model 分类设置独立阈值；
大规模数据集使用 DataLoader：给--max_data_loader_n_workers传入 ≥1 的值即可启用多进程图像预加载，配合--batch_size能大幅缩短整体打标耗时；
保持标签规范一致：若按 Animagine XL 3.1 等特定模型训练，建议完全复刻本文 3.1 节的参数组合（--remove_underscore+--use_rating_tags_as_last_tag+--character_tags_first+--character_tag_expand+--always_first_tags），并在大规模运行前先用--debug抽查少量图像的输出；
注意联动规则：设置--remove_underscore后，所有需要匹配的标签参数都不要再带下划线；设置自定义--caption_separator后，undesired_tags与always_first_tags也要用同样的分隔符传入。

至此，你已经掌握了 sd-scripts 中 WD14Tagger 打标脚本的完整用法：从环境安装、模型自动下载，到按 Animagine XL 3.1 规范生成标签，再到分类阈值、标签替换、排序策略等全部参数，以及标签文件向训练元数据的衔接方式。这套流程可以直接用于 LoRA、DreamBooth、Textual Inversion 等训练任务的数据集准备阶段。