背景
随着AI生成图像技术的迅猛发展，特别是生成对抗网络（GANs）和深度学习的不断进步，生成的图像变得越来越逼真。
这项技术不仅催生了许多创新应用，也带来了潜在的风险和挑战。
Transformer模型在图像识别中的作用
1、特征学习能力：Transformer模型具有强大的特征提取和表示能力，能够从图像中学习到细微的特征差异，识别出AI生成图像与真实图像之间的细微区别。
2、上下文理解：与传统的卷积神经网络（CNN）相比，Transformer模型更擅长捕捉图像中的全局上下文信息，使其在识别AI生成图像时，尤其在细节和纹理方面，表现得更加精准。
3、适应性强：通过预训练和微调，Transformer模型能够适应各种图像生成技术，保持高效的识别能力，即便面对不断进化的AI生成技术。
以下是一个利用Transformer模型来识别图像是否由AI生成的例子。
数据集
| 类别 | 训练集数量 | 测试集数量 |
| FAKE | 50,000 | 10,000 |
| REAL | 50,000 | 10,000 |
| 总计 | 100,000 | 20,000 |
完整步骤
1、导入包
import numpy as np
from datasets import load_dataset
import torch
from transformers importViTFeatureExtractor
from transformers importTrainingArguments
from transformers importTrainer
import tensorflow as tf
from tensorflow.keras.preprocessing.image importImageDataGenerator
from transformers importViTForImageClassification, default_data_collator
from torch.utils.data importDataLoader,Dataset
import os
from PIL import Image
2、图像进行预处理
from torch.utils.data importDataset
from PIL importImage
import os
import torch
classCustomImageDataset(Dataset):
def __init__(self, img_dir, feature_extractor):
"""
初始化数据集。
:param img_dir: 包含图像的根目录路径。应包含两个子目录：'REAL' 和 'FAKE'。
:param feature_extractor: 用于图像预处理的特征提取器。
"""
self.img_dir = img_dir
self.img_labels =[]#存储图像标签的列表
self.img_files =[]#存储图像文件路径的列表
self.feature_extractor = feature_extractor # 特征提取器
self.label_mapping ={'REAL':1,'FAKE':0}#标签映射字典
# 遍历 'REAL' 和 'FAKE' 目录
for label_dir in['REAL','FAKE']:
dir_path = os.path.join(img_dir, label_dir)#获取每个子目录的路径
files = os.listdir(dir_path)#获取目录中的所有文件
for file in files:
# 将每个文件的完整路径添加到 img_files 列表
self.img_files.append(os.path.join(dir_path, file))
# 将每个文件的标签添加到 img_labels 列表
self.img_labels.append(self.label_mapping[label_dir])
def __len__(self):
"""
返回数据集中图像的数量。
:return: 数据集中图像的总数。
"""
return len(self.img_files)
def __getitem__(self, idx):
"""
根据给定的索引返回图像和标签。
:param idx: 图像的索引。
:return: 一个字典，包含图像的张量 'pixel_values' 和标签的张量 'labels'。
"""
img_path =self.img_files[idx]#获取图像路径
image =Image.open(img_path).convert("RGB")#打开图像并转换为 RGB 格式
label =self.img_labels[idx]#获取对应的标签
features =self.feature_extractor(images=image, return_tensors="pt")#使用特征提取器处理图像
# 确保输出是字典格式，并移除多余的维度
return{"pixel_values": features['pixel_values'].squeeze(),"labels": torch.tensor(label)}
3、加载了一个本地保存的Vision Transformer模型和特征提取器，设置设备（GPU 或 CPU），并创建了用于训练和测试的数据集和数据加载器。
from transformers importViTFeatureExtractor,ViTForImageClassification
from torch.utils.data importDataLoader
import torch
# Vision Transformer (ViT) 模型在 ImageNet-21k（1400 万张图像，21,843 个类别）上以 224x224 分辨率进行预训练。
model_id ='google/vit-base-patch16-224-in21k'
# 定义本地模型文件路径
model_path ='../model'
# 加载特征提取器
feature_extractor =ViTFeatureExtractor.from_pretrained(
model_path,#使用本地路径
local_files_only=True# 仅使用本地文件
)
# 设定设备（GPU 或 CPU）
device = torch.device('cuda'if torch.cuda.is_available()else'cpu')
# 加载分类器模型
model =ViTForImageClassification.from_pretrained(
model_path,#使用本地路径
num_labels=2,#设置分类标签的数量
local_files_only=True# 仅使用本地文件
)
# 将模型移到所选设备（GPU 或 CPU）
model.to(device)
# 创建数据集
train_dataset =CustomImageDataset(img_dir='../dataset/train', feature_extractor=feature_extractor)
test_dataset =CustomImageDataset(img_dir='../dataset/test', feature_extractor=feature_extractor)
# 创建数据加载器
train_loader =DataLoader(train_dataset, batch_size=16, shuffle=True)#训练集的数据加载器
test_loader =DataLoader(test_dataset, batch_size=16, shuffle=False) # 测试集的数据加载器
4、打印当前设备，查看使用的设备类型。
# 打印当前设备
print("当前设备:", device)
5、设置了训练模型的参数，并使用这些参数初始化了一个Trainer实例
from transformers importTrainer,TrainingArguments, default_data_collator
# 配置训练参数
training_args =TrainingArguments(
per_device_train_batch_size=4,#每个设备上训练的批次大小为4
evaluation_strategy="epoch",#每个训练周期结束时进行评估
num_train_epochs=4,#总共训练4个周期（epochs）
save_strategy="epoch",#每个训练周期结束时保存模型
logging_steps=10,#每10个训练步骤记录一次日志
learning_rate=2e-4,#学习率设置为0.0002
save_total_limit=2,#保留最新的2个模型检查点，删除旧的检查点
remove_unused_columns=False,#不移除数据集中未使用的列
push_to_hub=False,#不将模型推送到HuggingFaceHub
load_best_model_at_end=True,#在训练结束时加载性能最佳的模型
output_dir="./outputs",#模型和其他输出保存到指定目录
use_cpu=False# 不强制使用 CPU，默认使用 GPU（如果可用）
)
# 初始化 Trainer
trainer =Trainer(
model=model,#使用之前定义的模型
args=training_args,#使用上面定义的训练参数
train_dataset=train_dataset,#训练数据集
eval_dataset=test_dataset,#评估数据集
data_collator=default_data_collator,#数据整理器，用于处理数据批次
compute_metrics=None,#计算指标的函数，此处不计算任何指标
)
6、开始训练
- trainer.train()
时间很长。
7、将trainer实例中的模型（包括模型的权重和配置）保存到指定的目录。
- trainer.save_model("./outputs/final model")
8、开始验证
from transformers importAutoFeatureExtractor,AutoModelForImageClassification
import os
import torch
from PIL importImage
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif']=['Microsoft YaHei']
plt.rcParams['axes.unicode_minus']=False
# 使用绝对路径
model_path = os.path.abspath("D:\MY\8-m\final model")
# 尝试加载特征提取器和模型
try:
feature_extractor =AutoFeatureExtractor.from_pretrained(model_path, local_files_only=True)
exceptOSErroras e:
print(f"加载特征提取器时出错: {e}")
try:
model =AutoModelForImageClassification.from_pretrained(model_path, local_files_only=True)
exceptOSErroras e:
print(f"载模型时出错: {e}")
# 设备设置
device = torch.device("cuda"if torch.cuda.is_available()else"cpu")
model.to(device)
# 加载和预处理图像
image_path ='潘展乐.png'
image =Image.open(image_path).convert("RGB")
inputs = feature_extractor(images=image, return_tensors="pt")
pixel_values = inputs['pixel_values'].to(device)
# 执行预测
model.eval()
with torch.no_grad():
outputs = model(pixel_values)
# 解析预测结果
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
predicted_label ='真实'if predicted_class_idx ==1else'AI生成'
# 打印预测标签
print(f"预测标签: {predicted_label}")
# 展示图像和预测结果
plt.imshow(image)
plt.axis('off')
plt.title(f"预测标签: {predicted_label}")
plt.show()
over，模型的准确率还是很棒的。
