自训Transformer模型：识别图像是否由AI生成？

背景

随着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

50,000个真实图片和50,000个AI生成图片

完整步骤

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()