项目简介
JoyTag是一个基于深度学习的图像标注工具，它可以自动识别图片中的内容并生成相关标签。这个项目利用先进的机器学习算法，不仅提高了标记的准确性，还大幅度提升了处理速度，适用于大规模图像数据集的快速标注。JoyTag的开源特性允许开发者自由使用和修改代码，支持自定义训练模型以适应特定的标注需求，使其在图像识别和自动标注领域具有广泛的应用前景。
扫码加入交流群
获得更多技术支持和交流
使用模型
from Models import VisionModelfrom PIL import Imageimport torch.amp.autocast_modefrom pathlib import Pathimport torchimport torchvision.transforms.functional as TVFpath = '/home/.../joytag/models' # Change this to where you downloaded the modelTHRESHOLD = 0.4model = VisionModel.load_model(path)model.eval()model = model.to('cuda')with open(Path(path) / 'top_tags.txt', 'r') as f: top_tags = [line.strip() for line in f.readlines() if line.strip()]def prepare_image(image: Image.Image, target_size: int) -> torch.Tensor: # Pad image to square image_shape = image.size max_dim = max(image_shape) pad_left = (max_dim - image_shape[0]) // 2 pad_top = (max_dim - image_shape[1]) // 2 padded_image = Image.new('RGB', (max_dim, max_dim), (255, 255, 255)) padded_image.paste(image, (pad_left, pad_top)) # Resize imageif max_dim != target_size: padded_image = padded_image.resize((target_size, target_size), Image.BICUBIC) # Convert to tensor image_tensor = TVF.pil_to_tensor(padded_image) / 255.0 # Normalize image_tensor = TVF.normalize(image_tensor, mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])return image_tensor@torch.no_grad()def predict(image: Image.Image): image_tensor = prepare_image(image, model.image_size) batch = {'image': image_tensor.unsqueeze(0).to('cuda'), } with torch.amp.autocast_mode.autocast('cuda', enabled=True): preds = model(batch) tag_preds = preds['tags'].sigmoid().cpu() scores = {top_tags[i]: tag_preds[0][i] for i in range(len(top_tags))} predicted_tags = [tag for tag, score in scores.items() if score > THRESHOLD] tag_string = ', '.join(predicted_tags)return tag_string, scoresimage = Image.open('test.jpg')tag_string, scores = predict(image)print(tag_string)for tag, score in sorted(scores.items(), key=lambda x: x[1], reverse=True): print(f'{tag}: {score:.3f}')
目标
大多数公共视觉模型在其训练数据集中包含了大量的过滤工作。这意味着，今天的基础视觉模型在许多概念上的基本层面上表现较弱。这限制了表达自由、包容性和多样性。这也限制了机器学习模型对我们世界的基本理解及将来的理解能力。JoyTag团队认为，人类用户应该自由表达自己，不应该因为随意和反复无常的内容过滤而受到歧视。
JoyTag团队也相信，机器学习模型应该对世界有一个广泛、深入和包容的理解。这并不排除在训练后使用对齐技术来减少模型的偏见，但它确实排除了使用过滤或对齐，这些做法减少了模型理解世界的能力或用户表达自己的能力。
开发手册
模型目前主要受到数据的限制，如果超出当前的训练方案或模型大小，就会发生过拟合。基于L/14@224的模型能够训练到F1得分为0.51之前就开始过拟合。对于H/14也是如此，即使增加了StochDepth，F1得分只有小幅改善然后过拟合。这是预期之中的，因为B/16@448的计算成本与L/14@224相似，但参数更少，因此似乎提供了一种在保持良好正则化的同时扩展模型的方法。
在我在220M和440M规模的测试中，MixUp和标签平滑都没有提供改进。Mixup提供了最多的正则化，但最终F1得分比没有使用它时更低。当前的学习率、权重衰减和stochdepth率在220M范围内基于扫描似乎接近最优。对于448训练部分的设置没有进行实验，所以那里可能有更优的学习率。
简单增强比RandAugment和其他增强方案工作得更好，提供了对抗过拟合的最佳治疗。可能其他增强方式能提供更多的正则化并允许更多的训练。
当过拟合时，模型的F1倾向于继续改善，而验证损失下降。
在220M扫描中，焦点损失似乎与普通的BCE损失相比没有显著差异。按标签频率加权没有改善结果。
稳定训练必须使用损失乘数，可能是由于与典型BCE任务相比，标签的数量和极端分布导致损失过低，以至于GradScaler无法考虑。简单地将损失乘以128或256就足以稳定训练。
AdamW也能工作，并且可以提供非常小的F1改进，但LAMB结果在训练曲线更稳定，且不需要调整预热计划。
使用CNN主干而不是ViTs的传统打补丁主干，在B/16模型上大幅提升了F1。对L/14模型似乎帮助较小。更多的调查可能会有成果。
项目链接
https://github.com/fpgaminer/joytag
关注「 开源AI项目落地 」公众号
