微调视觉Transformer（ViT） 是一种Transformer编码器架构的变体，与BERT类似，已适配于图像分类任务。这款名为“google/vit-base-patch16-224-in21k”的特定模型，是在大量图像集合上通过监督方式预训练的，其训练数据来源于ImageNet-21k数据集。预训练数据集中的图像均调整为224x224像素的分辨率，使其适用于广泛的图像识别任务。
在训练阶段，为确保模型达到最佳性能，我们对超参数设置给予了细致关注。模型的微调采用了经过审慎选择的16 batch size。这一选择不仅平衡了计算效率，还使模型能够有效地处理和学习多样化的图像。
为助力微调过程，我们采用了5e-5的学习率。学习率作为关键的调整参数，决定了训练期间模型参数调整的幅度。在此案例中，选择5e-5的学习率旨在在快速收敛与稳定优化之间取得和谐平衡，从而得到一个不仅学习速度快，且能在整个训练过程中稳步提升自身能力的模型。
此训练阶段使用的是一个专有数据集，其中包含80,000张图像，每张图像都具有相当程度的变异性。该数据集经过精心策划，包含“normal”和“nsfw”两个不同类别。这种多样性使模型能够掌握细微的视觉模式，使其具备准确区分安全内容与不良内容的能力。
这一细致训练过程的总体目标是让模型深刻理解视觉线索，确保其在处理特定的NSFW图像分类任务时具备稳健性和能力。最终得到的模型已准备好为内容安全和审核做出重要贡献，同时保持最高标准的准确性和可靠性。
from openmind import pipeline, is_torch_npu_available
from openmind_hub import snapshot_download
from PIL import Image
import torch
import argparse
import time
import requests
from io import BytesIO
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_name_or_path",
type=str,
help="Path to model",
default="jeffding/nsfw_image_detection-openmind",
)
args = parser.parse_args()
return args
def main():
args = parse_args()
model_path = args.model_name_or_path
if is_torch_npu_available():
device = "npu:0"
else:
device = "cpu"
start_time = time.time()
yzmdata = requests.get("https://ankur3107.github.io/assets/images/image-captioning-example.png")
tempIm = BytesIO(yzmdata.content)
img = Image.open(tempIm)
classifier = pipeline("image-classification", model=model_path,device_map=device)
result = classifier(img)
print(result)
end_time = time.time()
print(f"硬件环境：{device},推理执行时间：{end_time - start_time}秒")
if __name__ == "__main__":
main()以下是如何使用此模型基于 2 个类别（normal、nsfw）中的 1 个对图像进行分类的方法：
# Use a pipeline as a high-level helper
from PIL import Image
from transformers import pipeline
img = Image.open("<path_to_image_file>")
classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
classifier(img)
# Load model directly
import torch
from PIL import Image
from transformers import AutoModelForImageClassification, ViTImageProcessor
img = Image.open("<path_to_image_file>")
model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection')
with torch.no_grad():
inputs = processor(images=img, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
predicted_label = logits.argmax(-1).item()
model.config.id2label[predicted_label]
该模型的训练数据包含一个专有数据集，由大约80,000张图像组成。该数据集具有显著的可变性，包含“normal”和“nsfw”两个不同类别。基于此数据的训练过程旨在使模型具备有效区分安全内容和不当内容的能力。
- 'eval_loss': 0.07463177293539047,
- 'eval_accuracy': 0.980375,
- 'eval_runtime': 304.9846,
- 'eval_samples_per_second': 52.462,
- 'eval_steps_per_second': 3.279
注意： 在实际应用中使用此模型时，必须以负责任和符合伦理的方式进行，遵守内容准则和适用法规，尤其是在涉及可能敏感内容的场景中。
有关模型微调与使用的更多详细信息，请参阅模型文档和模型中心。
免责声明： 模型的性能可能会受到其微调数据的质量和代表性的影响。建议用户评估该模型对其特定应用和数据集的适用性。