背景随着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、导入包importnumpyasnp
fromdatasetsimportload_dataset
importtorch
fromtransformersimportViTFeatureExtractor
fromtransformersimportTrainingArguments
fromtransformersimportTrainer
importtensorflowastf
fromtensorflow.keras.preprocessing.imageimportImageDataGenerator
fromtransformersimportViTForImageClassification,default_data_collator
fromtorch.utils.dataimportDataLoader,Dataset
importos
fromPILimportImage
2、图像进行预处理fromtorch.utils.dataimportDataset
fromPILimportImage
importos
importtorch
classCustomImageDataset(Dataset):
def__init__(self,img_dir,feature_extractor):
"""
初始化数据集。
:paramimg_dir:包含图像的根目录路径。应包含两个子目录：'REAL'和'FAKE'。
:paramfeature_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'目录
forlabel_dirin['REAL','FAKE']:
dir_path=os.path.join(img_dir,label_dir)#获取每个子目录的路径
files=os.listdir(dir_path)#获取目录中的所有文件
forfileinfiles:
#将每个文件的完整路径添加到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:数据集中图像的总数。
"""
returnlen(self.img_files)
def__getitem__(self,idx):
"""
根据给定的索引返回图像和标签。
:paramidx:图像的索引。
: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），并创建了用于训练和测试的数据集和数据加载器。fromtransformersimportViTFeatureExtractor,ViTForImageClassification
fromtorch.utils.dataimportDataLoader
importtorch
#VisionTransformer(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'iftorch.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实例fromtransformersimportTrainer,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/finalmodel")
8、开始验证fromtransformersimportAutoFeatureExtractor,AutoModelForImageClassification
importos
importtorch
fromPILimportImage
importmatplotlib.pyplotasplt
plt.rcParams['font.sans-serif']=['MicrosoftYaHei']
plt.rcParams['axes.unicode_minus']=False
#使用绝对路径
model_path=os.path.abspath("D:\\MY\\8-m\\finalmodel")
#尝试加载特征提取器和模型
try:
feature_extractor=AutoFeatureExtractor.from_pretrained(model_path,local_files_only=True)
exceptOSErrorase:
print(f"加载特征提取器时出错:{e}")
try:
model=AutoModelForImageClassification.from_pretrained(model_path,local_files_only=True)
exceptOSErrorase:
print(f"载模型时出错:{e}")
#设备设置
device=torch.device("cuda"iftorch.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()
withtorch.no_grad():
outputs=model(pixel_values)
#解析预测结果
logits=outputs.logits
predicted_class_idx=logits.argmax(-1).item()
predicted_label='真实'ifpredicted_class_idx==1else'AI生成'
#打印预测标签
print(f"预测标签:{predicted_label}")
#展示图像和预测结果
plt.imshow(image)
plt.axis('off')
plt.title(f"预测标签:{predicted_label}")
plt.show()
over，模型的准确率还是很棒的。
| 
