Camie Tagger

🚀 Camie Tagger

Camie Tagger是一款先进的深度学习模型，用于自动为动漫/漫画插图添加跨多个类别的相关标签。在20,116个样本的测试集上，针对70,527个可能的标签，该模型实现了58.1%的微F1分数（使用平衡阈值预设时，宏F1分数为31.5%）。该模型在单张3060显卡上完成训练。

🚀 快速开始

项目更新（2025年4月 - 5月）

很抱歉这段时间没有更新。我一直忙于大学的事务，并且在额外训练（1个轮次）时犯了个错误，只加载了模型权重，而没有加载优化器状态。这很重要，因为优化器状态会为每个参数维护一个经过调整的学习率。我只加载权重是因为训练时间变得更慢了，所以忽略了加载优化器状态的重要性。

好消息是，我很快会拿到一张RTX 3090显卡。这意味着我可以在GPU显存中完整训练整个模型（无需卸载），并利用3090提供的3倍内存带宽。这应该会大幅缩短训练时间！

我也在进行一些大语言模型（LLM）实验，但距离发布还为时尚早。

✨ 主要特性

• 高效训练：仅在单张RTX 3060 GPU（12GB显存）上完成训练。
• 快速收敛：在1,756,098个批次中，对7,024,392个样本（3.52个轮次）进行训练。
• 全面覆盖：涵盖7个类别（通用、角色、版权、艺术家、元信息、评级、年份）下的70,527个标签。
• 创新架构：采用两阶段预测模型，以EfficientNetV2为骨干网络，结合嵌入层和交叉注意力机制处理标签上下文。
• 模型规模：初始模型（2.14亿参数），精炼模型（4.24亿参数）。
• 用户友好界面：易于使用的应用程序，支持自定义阈值。

本项目证明，通过正确的优化技术，在消费级硬件上也可以训练出高质量的动漫图像标签模型。

🔑 关键亮点

• 多类别标签系统：处理通用标签、角色、版权（系列）、艺术家、元信息和内容评级。
• 高性能：在70,527个可能的标签上，微F1分数达到53.4%（宏F1分数为31.5%）。
• Windows兼容性：仅初始模型模式可在不支持Flash Attention的Windows系统上运行。
• Streamlit网页界面：用户友好的界面，支持上传和分析图像，还有标签收集游戏。
• 可调整阈值配置文件：提供微、宏、平衡、特定类别、高精度和高召回率等配置文件。
• 细粒度控制：可针对每个类别调整阈值，以平衡精确率和召回率。
• Safetensors和ONNX支持：原始的pickle文件可在/models目录中找到。
• EfficientNetV2 - L骨干网络：通过精炼嵌入层，显著提高了骨干网络的性能。

📊 性能说明

从宏分数来看，尤其是通用标签的分数，模型的性能似乎有点不尽如人意（请查看性能部分的表格）。然而，我发现该模型在预测这些标签时总体表现良好，在角色和版权标签方面表现出色。而且，需要预测的标签数量确实很多。

好消息是，有明显证据表明该模型训练不足，在训练轮次中各项指标持续提升：

训练进度（微F1与宏F1）：
第2轮：  微F1：0.595    宏F1：0.226  
第3轮：  微F1：0.606    宏F1：0.268（+4.2%）
第3.5轮：微F1：0.611    宏F1：0.290（仅0.5轮，+2.2%）

这是合理的，因为3.5个轮次的训练时间确实不算长。

微F1与宏F1解释

• 微F1：通过考虑每个标签实例的预测结果，全局计算指标。该指标受常见标签和样本较多类别的影响较大。
• 宏F1：独立计算每个标签的指标，然后取平均值。这使得稀有标签和常见标签具有相同的权重。

宏F1的显著提升（每轮+4%）表明，更长时间的训练将特别有利于稀有标签的识别，而微F1的提升速度正在放缓，因为常见标签已经被很好地学习。

未来训练计划

我计划继续训练该模型，以进一步提高性能，特别是针对稀有标签。然而，在我目前的硬件上，每个轮次的训练大约需要1.5 - 2周的夜间时间。

如果您想支持在完整数据集上进行进一步训练或我的未来项目，请考虑在此处支持我：(https://ko-fi.com/camais)。您的支持将直接有助于更长时间的训练和更好的模型！

完成这个项目后，我计划转向大语言模型，因为我有很多关于如何改进它们的想法。我将根据社区的关注情况更新这个模型。

📈 性能分析

整体性能

初始预测

类别	配置文件	阈值	微F1	宏F1
总体	微优化	0.326	0.611	0.290
	宏优化	0.201	0.534	0.331
	平衡	0.258	0.581	0.315
	高精度	0.500	0.497	0.163
	高召回率	0.120	0.308	0.260
艺术家	微优化	0.262	0.474	0.295
	宏优化	0.140	0.262	0.287
	平衡	0.258	0.474	0.298
	高精度	0.464	0.310	0.135
	高召回率	0.153	0.302	0.301
角色	微优化	0.294	0.749	0.444
	宏优化	0.161	0.608	0.517
	平衡	0.258	0.746	0.478
	高精度	0.500	0.655	0.268
	高召回率	0.100	0.336	0.386
版权	微优化	0.334	0.789	0.325
	宏优化	0.205	0.700	0.404
	平衡	0.258	0.763	0.377
	高精度	0.500	0.747	0.209
	高召回率	0.100	0.347	0.267
通用	微优化	0.322	0.607	0.180
	宏优化	0.225	0.537	0.210
	平衡	0.258	0.576	0.204
	高精度	0.500	0.482	0.095
	高召回率	0.124	0.301	0.161
元信息	微优化	0.330	0.601	0.134
	宏优化	0.209	0.487	0.143
	平衡	0.258	0.557	0.144
	高精度	0.500	0.458	0.081
	高召回率	0.120	0.309	0.103
评级	微优化	0.359	0.808	0.791
	宏优化	0.359	0.808	0.791
	平衡	0.258	0.779	0.768
	高精度	0.500	0.738	0.686
	高召回率	0.100	0.650	0.611
年份	微优化	0.266	0.332	0.285
	宏优化	0.258	0.331	0.286
	平衡	0.258	0.331	0.286
	高精度	0.302	0.308	0.251
	高召回率	0.213	0.304	0.279

精炼预测

类别	配置文件	阈值	微F1	宏F1
总体	微优化	0.326	0.613	0.295
	宏优化	0.193	0.546	0.338
	平衡	0.262	0.586	0.326
	高精度	0.500	0.499	0.173
	高召回率	0.120	0.310	0.262
艺术家	微优化	0.278	0.480	0.297
	宏优化	0.148	0.288	0.299
	平衡	0.262	0.483	0.311
	高精度	0.480	0.314	0.140
	高召回率	0.153	0.302	0.306
角色	微优化	0.302	0.757	0.460
	宏优化	0.157	0.591	0.524
	平衡	0.262	0.751	0.496
	高精度	0.500	0.669	0.286
	高召回率	0.100	0.331	0.386
版权	微优化	0.367	0.792	0.317
	宏优化	0.189	0.671	0.419
	平衡	0.262	0.767	0.392
	高精度	0.492	0.755	0.228
	高召回率	0.100	0.349	0.270
通用	微优化	0.326	0.608	0.181
	宏优化	0.237	0.553	0.215
	平衡	0.262	0.580	0.208
	高精度	0.500	0.484	0.100
	高召回率	0.124	0.303	0.165
元信息	微优化	0.330	0.602	0.127
	宏优化	0.197	0.468	0.145
	平衡	0.262	0.563	0.152
	高精度	0.500	0.453	0.087
	高召回率	0.120	0.305	0.107
评级	微优化	0.375	0.808	0.787
	宏优化	0.338	0.809	0.795
	平衡	0.262	0.784	0.773
	高精度	0.500	0.735	0.678
	高召回率	0.100	0.652	0.610
年份	微优化	0.266	0.332	0.292
	宏优化	0.258	0.331	0.293
	平衡	0.262	0.333	0.293
	高精度	0.306	0.306	0.255
	高召回率	0.209	0.301	0.275

该模型在角色识别（在26,968个标签上F1分数达到75.7%）、版权/系列检测（在5,364个标签上F1分数达到79.2%）和内容评级分类（在4个标签上F1分数达到80.8%）方面表现尤其出色。

初始预测与精炼预测性能对比

预测类型	微F1	宏F1	精确率	召回率
初始	0.611	0.290	0.610	0.613
精炼	0.613	0.295	0.617	0.609

精炼过程使微F1提高了0.2%，宏F1提高了0.5%。如图所示，精炼预测相比初始预测有小幅但持续的提升，这使得仅初始模型成为不支持Flash Attention的Windows用户的不错选择。

实际标签准确率

在个人测试中，我发现根据基准测试被判定为“误报”的许多标签，实际上是Danbooru数据集中缺失的正确标签（该数据集本身的标签也并非100%完美）。一些观察结果如下：

• 对于角色、版权和艺术家类别，即使模型不太确定，预测的顶级标签也经常是正确的。
• 许多看似错误的通用标签实际上是合适的描述符，只是原始标签中未包含。

出于这些原因，尽管“高召回率”阈值配置文件的正式F1分数较低，但在实际应用中可能会产生更好的效果。使用该应用程序时，将每个类别输出限制为前N个标签可以获得最准确和有用的结果。

与WDTaggers的比较

比较我们的标签器有点棘手，原因有两个：

• 我的标签器有70,000多个标签，而WD标签器有10,000多个标签。
• 我只保留了至少有25个通用标签（以及至少1个角色 + 版权标签）的样本，而WD标签器至少保留10个标签的样本。在我自己的测试中，我的样本平均每个样本有44.74个标签，而WD标签器为36.48个，相差18.5%。

考虑到这些因素，对于当前的检查点，我认为我的标签器在稀有标签方面似乎更准确，因为它涵盖的标签更多。而WD标签器在常见标签方面明显更准确。

以下是两个标签器在相同标签数量下的比较：https://huggingface.co/datasets/SmilingWolf/camie-tagger-vs-wd-tagger-val

在我个人的测试中，我的标签器在识别稀有角色、版权、艺术家和一些稀有通用标签（如替代服装和艺术家等）方面似乎更好。不过，Camie标签器在通用标签方面似乎有1 - 3个更多的误报。请记住，这只是基于几张图像的测试，可能存在误差。

经过更多轮次的训练后，我认为差距会缩小很多，但这需要大约一个月的时间，因为我只能在夜间进行训练。我认为WD标签器训练了50多个轮次，而我的目前只训练了3.5个轮次。

总体而言，标签的分布呈现出极长的尾部。我的游戏展示了标签的稀有度范围，大多数标签都属于最稀有类别。两个标签器在常见标签上都应该能提供较好的准确性。

🛠️ 环境要求

• 特定版本的Python 3.11.9（更高版本不兼容）
• PyTorch 1.10+
• Streamlit
• PIL/Pillow
• NumPy
• Flash Attention（注意：在Windows上无法正常工作，仅精炼模型需要，且我对此支持有限）

🔧 使用方法

通过执行setup.bat脚本设置应用程序和游戏，该脚本将安装所需的虚拟环境：

• 上传您自己的图像或选择示例图像。
• 选择不同的阈值配置文件。
• 调整特定类别的阈值。
• 查看按类别组织的预测结果。
• 根据置信度过滤和排序标签。

使用run_app.bat和run_game.bat脚本运行应用程序和游戏。

🎮 标签收集游戏（Camie Collector）

这是一款标签收集游戏，以游戏化的方式进行动漫图像标签标注，帮助您了解模型的性能和局限性。这原本只是一个突发奇想，没想到发展成了一个项目！其中有很多Project Moon的元素。

游戏玩法

1. 上传一张图像。
2. 扫描标签以发现它们。
3. 发现新标签可获得标签币。
4. 使用标签币进行升级，降低阈值。
5. 降低阈值可揭示更稀有的标签！
6. 收集相关标签集可获得奖励，并揭示独特的马赛克图案！
7. 访问图书馆系统以发现独特的标签（不可收集）。
8. 使用收集的标签激发新的搜索或生成标签精华。
9. 使用脑啡肽生成标签精华。
10. 使用标签精华生成器收集标签及其相关标签。例如灯的精华：

🖥️ 网页界面指南

该界面分为三个主要部分：

1. 模型选择（侧边栏）：
   • 选择完整模型、仅初始模型或ONNX加速（仅初始模型）。
   • 查看模型信息和内存使用情况。
2. 图像上传（左面板）：
   • 上传您自己的图像或选择示例图像。
   • 查看所选图像。
3. 标签控制（右面板）：
   • 选择阈值配置文件。
   • 调整阈值以平衡精确率和召回率，以及微/宏指标。
   • 配置显示选项。
   • 查看按类别组织的预测结果。

显示选项

• 显示所有标签：显示所有标签，包括低于阈值的标签。
• 紧凑视图：隐藏进度条，使显示更简洁。
• 最小置信度：过滤掉置信度较低的预测结果。
• 类别选择：选择要包含在摘要中的类别。

界面截图

🧠 训练细节

数据集

该模型在经过精心筛选的Danbooru 2024数据集子集上进行训练，该数据集包含大量带有详细标签的动漫/漫画插图。

筛选过程

数据集按照以下约束条件进行筛选：

# 每个图像每个类别所需的最小标签数
min_tag_counts = {
    'general': 25, 
    'character': 1, 
    'copyright': 1, 
    'artist': 0, 
    'meta': 0
}

# 每个标签所需的最小样本数
min_tag_samples = {
    'general': 20, 
    'character': 40, 
    'copyright': 50, 
    'artist': 200, 
    'meta': 50
}

此筛选过程：

1. 首先移除样本数少于min_tag_samples指定数量的标签。
2. 然后移除每个类别标签数量不足的图像（如min_tag_counts所指定）。

训练数据

• 初始数据集大小：约3,000,000张经过筛选的图像。
• 训练子集：2,000,000张图像（由于存储和时间限制）。
• 训练轮次：3.5个轮次。

预处理

图像经过最少的变换进行预处理：

• 张量归一化（缩放到0 - 1范围）。
• 调整大小并保持原始宽高比。
• 未应用额外的增强操作。

损失函数

该模型采用专门的UnifiedFocalLoss损失函数，以应对多标签预测中固有的极端类别不平衡问题：

class UnifiedFocalLoss(nn.Module):
    def __init__(self, device=None, gamma=2.0, alpha=0.25, lambda_initial=0.4):
        # 实现细节...

关键组件

1. 焦点损失机制：
   • 降低分类良好的样本权重（γ = 2.0），使训练集中在难分类的标签上。
   • 解决正样本和负样本之间的极端不平衡问题（通常为100:1或更糟）。
   • 使用α = 0.25来平衡70,527个可能标签中的正/负样本。
2. 两阶段加权：
   • 结合两个预测阶段（initial_predictions和refined_predictions）的损失。
   • 使用λ = 0.4对初始预测损失进行加权，给予精炼预测更多的权重（0.6）。
   • 这鼓励模型在精炼阶段改进预测，同时保持较强的初始预测。
3. 每个样本统计信息：
   • 分别跟踪正样本和负样本的指标。
   • 提供关于预测分布的详细调试信息。
   • 能够分析哪些标签类别表现良好或不佳。

尽管要面对70,527个可能标签的极端分类挑战，但这个损失函数对于在不同标签类别上实现高F1分数至关重要。

DeepSpeed配置

Microsoft DeepSpeed对于在消费级硬件上训练该模型至关重要。该项目使用经过精心调整的配置以最大化效率：

def create_deepspeed_config(
    config_path,
    learning_rate=3e-4,
    weight_decay=0.01,
    num_train_samples=None,
    micro_batch_size=4,
    grad_accum_steps=8
):
    # 实现细节...

关键优化

1. 内存效率：
   • ZeRO Stage 2：对优化器状态和梯度进行分区，显著减少内存需求。
   • 激活检查点：通过在反向传播期间重新计算激活值，以计算换取内存。
   • 连续内存优化：减少内存碎片。
2. 混合精度训练：
   • FP16模式：大多数计算使用半精度（16位），并自动进行损失缩放。
   • 初始缩放幂：设置为16，以确保在大批次大小下稳定收敛。
3. 梯度累积：
   • 微批次大小为4，梯度累积步数为8。
   • 有效批次大小为32，而一次只需要4个样本的内存。
4. 学习率调度：
   • 使用WarmupLR调度器，从3e - 6逐渐增加到3e - 4。
   • 在1/4个轮次内进行预热，以稳定早期训练。

这个配置使模型能够在仅12GB显存的情况下高效训练，同时在数百万个训练样本和70,527个输出维度上保持数值稳定性。

模型架构

该模型采用新颖的两阶段预测方法，与传统的单阶段模型相比，性能更优：

图像特征提取

• 骨干网络：EfficientNet V2 - L从输入图像中提取高质量的视觉特征。
• 空间池化：自适应平均池化将空间特征转换为紧凑的1280维嵌入。

初始预测阶段

• 通过多层分类器直接从图像特征进行分类。
• 瓶颈架构，线性层之间使用LayerNorm和GELU激活函数。
• 输出所有70,527个可能标签的初始概率。
• 模型大小：2.14657273亿个参数。

标签上下文机制

• 预测的顶级标签使用共享嵌入空间进行嵌入。
• 自注意力层允许标签根据共现模式相互影响。
• 归一化的标签嵌入表示图像的连贯“标签上下文”。

交叉注意力精炼

• 图像特征和标签嵌入通过交叉注意力进行交互。
• 图像特征的每个维度关注标签空间中的相关维度。
• 这在视觉特征和语义标签之间创建了双向信息流。

精炼预测

• 融合特征（原始特征 + 交叉注意力特征）输入到最终分类器。
• 残差连接确保在有益时保留初始预测。
• 温度缩放提供校准后的概率输出。
• 总模型大小：4.2479372亿个参数。

这种两阶段方法使模型能够利用标签共现模式和语义关系，在不显著增加参数数量的情况下提高准确性。

模型详情

标签类别

该模型识别以下类别的标签：

• 通用：视觉元素、概念、服装等（30,841个标签）。
• 角色：图像中出现的单个角色（26,968个标签）。
• 版权：源材料（动漫、漫画、游戏）（5,364个标签）。
• 艺术家：作品的创作者（7,007个标签）。
• 元信息：关于图像的元信息（323个标签）。
• 评级：内容评级（4个标签）。
• 年份：上传年份（20个标签）。

所有支持的标签存储在model/metadata.json文件中，该文件将标签ID映射到其名称和类别。

💻 训练环境

该模型在令人惊讶的适度硬件上进行训练：

• GPU：单张NVIDIA RTX 3060（12GB显存）。
• 内存：64GB系统内存。
• 平台：带有WSL（Windows Subsystem for Linux）的Windows系统。
• 库：
  • Microsoft DeepSpeed用于内存高效训练。
  • 支持CUDA加速的PyTorch。
  • Flash Attention用于优化注意力计算。

训练笔记本

该仓库包含两个主要的训练笔记本：

1. CAMIE Tagger.ipynb：
   • 主要的训练笔记本。
   • 数据集加载和预处理。
   • 模型初始化。
   • 集成DeepSpeed的初始训练循环。
   • 标签选择优化。
   • 指标跟踪和可视化。
2. Camie Tagger Cont and Evals.ipynb：
   • 从检查点继续训练。
   • 全面的模型评估。
   • 每个类别的性能指标。
   • 阈值优化。
   • 模型转换，以便在应用程序中部署。
   • 独立应用程序的导出功能。

训练监控器

该项目包含一个实时训练监控器，训练期间可通过浏览器在localhost:5000访问：

性能提示

⚠️ 重要提示

为了获得最佳训练速度，请将VSCode最小化，并在浏览器中打开训练监控器。由于Windows/WSL图形堆栈处理窗口焦点和CUDA内核执行的方式，这可以将迭代速度提高3 - 5倍。

监控器功能

训练监控器提供三个主要视图：

1. 概览页面

• 训练进度：实时指标，包括轮次、批次、速度和时间估计。
• 损失图表：训练和验证损失可视化。
• F1分数：训练和验证的初始和精炼F1指标。

2. 预测页面

• 图像预览：显示当前正在分析的样本。
• 预测控制：在初始和精炼预测之间切换。
• 标签分析：
  • 颜色编码的标签结果（正确、错误、缺失）。
  • 置信度可视化，带有概率条。
  • 按类别组织。
  • 用于错误分析的过滤选项。

3. 选择分析页面

• 选择指标：标签选择质量的统计信息。
  • 真实标签召回率。
  • 真实标签与非真实标签的平均概率。
  • 选择的唯一标签。
• 选择图表：选择质量随时间的趋势。
• 所选标签详情：模型选择的标签及其置信度分数的详细视图。

该监控器为两阶段预测模型的性能提供了宝贵的见解，特别是初始和精炼预测阶段之间的标签选择过程。

训练注意事项

• 训练笔记本需要WSL，并且可能需要32GB以上的内存来处理数据集。
• Microsoft DeepSpeed对于将模型和批次放入可用显存至关重要。
• 尽管硬件有限，但模型取得了令人印象深刻的结果。
• 如果有更多的计算资源，该模型可以在完整数据集上进行更长时间的训练。

🙏 致谢

• Claude Sonnet 3.5和3.7在头脑风暴和编码方面提供了极大的帮助。
• EfficientNetV2作为骨干网络。
• Danbooru提供了令人惊叹的带标签动漫图像数据集。
• p1atdev提供了处理后的Danbooru 2024数据集。
• Microsoft的DeepSpeed使在消费级硬件上进行训练成为可能。
• PyTorch和开源机器学习社区。

📄 许可证

本项目采用GPL - 3.0许可证。