ai-source-detector 是一个基于 ViT-Base 架构（86M 参数）的多分类图像分类模型，用于检测图像是由何种 AI 工具生成。模型能够区分五种来源：Stable Diffusion、Midjourney、DALL-E、真实照片以及其他 AI 工具生成的图像。适用于 AI 生成图像的溯源分析场景。
image-classification（图像五分类：stable_diffusion / midjourney / dalle / real / other_ai）
ONNX（CPU）/ PyTorch（NPU）
ViTForImageClassification（ViT-Base，86M 参数）
该模型原始格式为 ONNX，在 CPU 上使用 ONNX Runtime 进行推理。在 NPU 适配中，采用 onnx2torch 将 ONNX 模型转换为 PyTorch 格式，然后在昇腾 NPU（Atlas 800 A2/A3）上使用 torch_npu 进行推理。需要注意的是，CPU 和 NPU 使用不同的推理框架（ONNX Runtime vs torch_npu），但转换后的模型精度与原始 ONNX 模型高度一致。
# 基础环境
pip install torch torchvision
pip install onnxruntime
pip install onnx2torch
pip install Pillow numpy
# NPU 环境（需在昇腾设备上）
pip install torch_npu# CPU 推理（ONNX Runtime）
python accuracy_run.py --model_path ./model_files/model.onnx --image_path ./test_images/test.jpg --device cpu --framework onnx
# NPU 推理（torch_npu）
python accuracy_run.py --model_path ./model_files/model.onnx --image_path ./test_images/test.jpg --device npu --framework torchpython accuracy_run.py --model_path ./model_files/model.onnx --image_path ./test_images/test.jpg --device cpu --framework onnx --save_output cpu_output.npy
python accuracy_run.py --model_path ./model_files/model.onnx --image_path ./test_images/test.jpg --device npu --framework torch --save_output npu_output.npy
python compare_outputs.py --cpu_output cpu_output.npy --npu_output npu_output.npy| 平台 | 预测结果 | logits | probabilities | 推理耗时 |
|---|---|---|---|---|
| CPU | real | [0.100, -0.010, -0.440, 0.319, -0.508] | [0.234, 0.210, 0.137, 0.292, 0.128] | 135 ms |
| NPU | real | [0.100, -0.010, -0.440, 0.319, -0.508] | [0.234, 0.210, 0.137, 0.292, 0.128] | 250 ms |
| 指标 | 值 |
|---|---|
| Logits MAE | 0.000064 |
| Logits 相对误差 | 0.12% |
| Prob MAE | 0.000012 |
| 类别匹配率 | 100% |
| 分类结果一致性 | 完全一致 |
NPU 与 CPU 推理结果误差 < 1%
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
# 图像预处理
preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
def load_and_preprocess_image(image_path):
img = Image.open(image_path).convert('RGB')
input_tensor = preprocess(img).unsqueeze(0)
return input_tensor
# CPU 推理（ONNX Runtime）
import onnxruntime
def cpu_inference(model_path, input_tensor):
ort_session = onnxruntime.InferenceSession(model_path)
inputs = {ort_session.get_inputs()[0].name: input_tensor.numpy()}
outputs = ort_session.run(None, inputs)
return torch.from_numpy(outputs[0])
# NPU 推理（onnx2torch + torch_npu）
import torch_npu
from onnx2torch import convert
def npu_inference(model_path, input_tensor):
torch_model = convert(model_path)
torch_model = torch_model.to('npu').eval()
input_tensor = input_tensor.to('npu')
with torch.no_grad():
outputs = torch_model(input_tensor)
return outputs.cpu()
# 后处理
CLASS_NAMES = ['stable_diffusion', 'midjourney', 'dalle', 'real', 'other_ai']
def get_prediction(logits):
probs = torch.softmax(logits, dim=-1)
pred_class = torch.argmax(probs, dim=-1).item()
return CLASS_NAMES[pred_class], probs.numpy()| 平台 | 推理耗时 | 速度提升 |
|---|---|---|
| CPU | 135 ms | 1.0x（基线） |
| NPU | 250 ms | 0.54x |
注：该模型在 CPU 上使用 ONNX Runtime 进行推理，在 NPU 上使用 onnx2torch 转换后的 PyTorch 模型进行推理。NPU 推理耗时包含 ONNX 到 PyTorch 模型转换的额外开销以及 NPU 数据传输开销，同时 ONNX Runtime 在 CPU 上的推理优化较好，导致 NPU 在该模型上未体现出速度优势。对于追求性能的场景，建议进一步优化模型转换流程或使用昇腾 CANN 直接加载 ONNX 模型。
| 文件 | 说明 |
|---|---|
| model_files/model.onnx | ONNX 格式模型文件 |
| accuracy_run.py | 精度验证脚本 |
| test_images/ | 测试图像目录 |
| README.md | 本文件 |
本仓库提供完整的推理脚本，支持 CPU 和 NPU 双平台推理：
# NPU 推理
python3 inference.py --device npu
# CPU 推理
python3 inference.py --device cpu推理完成后会输出推理结果和耗时，表明模型在 NPU 上推理成功。
#+NPU #+CV #+图像分类 #+昇腾 #+AI检测 #+多分类 #+ViT #+ONNX