标签:

• 图像分类
• timm 库名称: timm 许可证: apache-2.0

nsfw-image-detection-384 模型卡

注意：与所有模型一样，此模型也可能出错。NSFW（不适合工作场所）内容具有主观性和情境性，本模型旨在帮助识别此类内容，使用风险自负。

Marqo/nsfw-image-detection-384 是一款轻量级图像分类模型，专为识别NSFW图像而设计。该模型体积约为其他开源模型的1/18至1/20，并在我们的数据集上实现了98.56%的优异准确率。该模型使用384x384像素的输入图像，并采用16x16像素的图块处理。

此模型基于包含22万张图片的专有数据集训练而成。训练集包含10万张NSFW样本和10万张SFW（适合工作场所）样本，测试集则包含1万张NSFW样本和1万张SFW样本。该数据集涵盖多样化的内容类型，包括：真实照片、绘画、Rule 34素材、表情包以及AI生成图像。NSFW的定义可能因情境而异，我们的数据集虽包含具有挑战性的样本，但可能无法100%契合所有使用场景，因此建议通过实验调整阈值以确定模型是否满足您的需求。

模型使用

使用timm进行图像分类

pip install timm

from urllib.request import urlopen
from PIL import Image
import timm
import torch

img = Image.open(urlopen(
    'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
))

model = timm.create_model("hf_hub:Marqo/nsfw-image-detection-384", pretrained=True)
model = model.eval()

data_config = timm.data.resolve_model_data_config(model)
transforms = timm.data.create_transform(**data_config, is_training=False)

with torch.no_grad():
    output = model(transforms(img).unsqueeze(0)).softmax(dim=-1).cpu()

class_names = model.pretrained_cfg["label_names"]
print("概率:", output[0])
print("类别:", class_names[output[0].argmax()])

评估

该模型在我们的数据集上表现优于现有NSFW检测器，以下是与其他模型的对比评估： AdamCodd/vit-base-nsfw-detector 和 Falconsai/nsfw_image_detection:

阈值与精确率/召回率权衡

调整NSFW概率阈值可在精确率、召回率和准确率之间进行权衡。这在需要不同置信度的应用场景中可能非常有用。

训练细节

本模型基于timm/vit_tiny_patch16_384.augreg_in21k_ft_in1k微调而成。

训练参数

batch_size: 256
color_jitter: 0.2
color_jitter_prob: 0.05
cutmix: 0.1
drop: 0.1
drop_path: 0.05
epoch_repeats: 0.0
epochs: 20
gaussian_blur_prob: 0.005
hflip: 0.5
lr: 5.0e-05
mixup: 0.1
mixup_mode: batch
mixup_prob: 1.0
mixup_switch_prob: 0.5
momentum: 0.9
num_classes: 2
opt: adamw
remode: pixel
reprob: 0.5
sched: cosine
smoothing: 0.1
warmup_epochs: 2
warmup_lr: 1.0e-05
warmup_prefix: false

引用文献

@article{dosovitskiy2020vit,
  title={An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale},
  author={Dosovitskiy, Alexey and Beyer, Lucas and Kolesnikov, Alexander and Weissenborn, Dirk and Zhai, Xiaohua and Unterthiner, Thomas and  Dehghani, Mostafa and Minderer, Matthias and Heigold, Georg and Gelly, Sylvain and Uszkoreit, Jakob and Houlsby, Neil},
  journal={ICLR},
  year={2021}
}

@misc{rw2019timm,
  author = {Ross Wightman},
  title = {PyTorch Image Models},
  year = {2019},
  publisher = {GitHub},
  journal = {GitHub repository},
  doi = {10.5281/zenodo.4414861},
  howpublished = {\url{https://github.com/huggingface/pytorch-image-models}}
}