WD ViT Tagger v3

支持评分、角色和通用标签。

使用 https://github.com/SmilingWolf/JAX-CV 进行训练。
训练所使用的 TPUs 由 TRC program 友情提供。

数据集

最后图片 ID：7220105
使用 Danbooru 图片 ID 取模 0000-0899 进行训练。
使用 ID 取模 0950-0999 的图片进行验证。
过滤掉了通用标签少于 10 张的图片。
过滤掉了图片数量少于 600 张的标签。

验证结果

v2.0: P=R: threshold = 0.2614, F1 = 0.4402
v1.0: P=R: threshold = 0.2547, F1 = 0.4278

更新内容

Model v2.0/Dataset v3：
训练了更多轮次。
使用基于标签频率的损失缩放来对抗类别不平衡问题。

Model v1.1/Dataset v3：
修改了 JAX 模型配置文件：添加了图片尺寸。
训练权重未做更改。

Model v1.0/Dataset v3：
更多训练图片，更多最新标签（截止至 2024-02-28）。
现在与 timm 兼容！加载并尝试使用标准的一行代码！
ONNX 模型与 v2 系列模型开发的代码兼容。
ONNX 模型的批次维度不再固定为 1。现在可以疯狂使用批次推理。
改用 Macro-F1 来评估模型性能，以便更好地了解整体训练进展。

运行时依赖

ONNX 模型需要 onnxruntime >= 1.17.0

推理代码示例

For timm: https://github.com/neggles/wdv3-timm
For ONNX: https://huggingface.co/spaces/SmilingWolf/wd-tagger
For JAX: https://github.com/SmilingWolf/wdv3-jax

结语

内容可能随时变更和更新。
建议下游用户使用标记版本，而非依赖仓库的最新提交。