（十）ps识别:Swin Transformer-T 与 ResNet50 结合的 PS 痕迹识别模型训练过程解析

Swin Transformer-T 与 ResNet50 结合的 PS 痕迹识别模型

思路分析

模型融合思路：
- 利用ResNet50提取图像的局部纹理和边缘特征，这对检测篡改区域的细微变化非常重要
- 利用Swin Transformer-T捕捉全局上下文信息和长距离依赖关系，有助于理解图像整体一致性
- 通过特征融合策略结合两种模型的输出，兼顾局部细节和全局语义
特征融合策略：
- 采用晚期融合策略，将两种模型的高层特征进行拼接
- 加入注意力机制，让模型自动学习不同特征的重要性权重
- 使用多层感知机(MLP)进行最终的分类决策
训练策略：
- 使用交叉熵损失函数处理二分类问题(真实/篡改)
- 采用学习率调度策略，动态调整训练过程
- 加入数据增强技术，提高模型泛化能力

代码实现

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from torchvision.transforms import functional as F
from timm import create_model
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import random# 设置随机种子，确保结果可复现
def set_seed(seed=42):random.seed(seed)np.random.seed(seed)torch.manual_seed(seed)if torch.cuda.is_available():torch.cuda.manual_seed(seed)torch.cuda.manual_seed_all(seed)torch.backends.cudnn.deterministic = Truetorch.backends.cudnn.benchmark = Falseset_seed()# 自定义数据集类
class PSTraceDataset(Dataset):def __init__(self, root_dir, transform=None, train=True):"""Args:root_dir (string): 数据集根目录，包含'original'和'tampered'两个子文件夹transform (callable, optional): 应用于样本的变换train (bool): 是否为训练集"""self.root_dir = root_dirself.transform = transformself.train = train# 加载原始图像和篡改图像的路径self.original_images = [os.path.join(root_dir, 'original', f) for f in os.listdir(os.path.join(root_dir, 'original')) if f.endswith(('png', 'jpg', 'jpeg'))]self.tampered_images = [os.path.join(root_dir, 'tampered', f) for f in os.listdir(os.path.join(root_dir, 'tampered')) if f.endswith(('png', 'jpg', 'jpeg'))]# 平衡数据集（取数量较少的类别作为基准）min_count = min(len(self.original_images), len(self.tampered_images))self.original_images = self.original_images[:min_count]self.tampered_images = self.tampered_images[:min_count]# 创建标签：0表示原始图像，1表示篡改图像self.images = self.original_images + self.tampered_imagesself.labels = [0] * min_count + [1] * min_count# 划分训练集和验证集（8:2）if self.train:split_idx = int(0.8 * len(self.images))self.images = self.images[:split_idx]self.labels = self.labels[:split_idx]else:split_idx = int(0.8 * len(self.images))self.images = self.images[split_idx:]self.labels = self.labels[split_idx:]def __len__(self):return len(self.images)def __getitem__(self, idx):img_path = self.images[idx]image = Image.open(img_path).convert('RGB')label = self.labels[idx]if self.transform:image = self.transform(image)return image, label# 定义数据增强和预处理
def get_transforms():train_transform = transforms.Compose([transforms.Resize((224, 224)),transforms.RandomHorizontalFlip(p=0.5),transforms.RandomVerticalFlip(p=0.5),transforms.RandomRotation(degrees=15),transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])val_transform = transforms.Compose([transforms.Resize((224, 224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])return train_transform, val_transform# 定义融合模型
class PSTraceDetector(nn.Module):def __init__(self, num_classes=2):super(PSTraceDetector, self).__init__()# 加载预训练的ResNet50self.resnet = models.resnet50(pretrained=True)# 移除最后的全连接层，保留特征提取部分self.resnet_features = nn.Sequential(*list(self.resnet.children())[:-1])# 加载预训练的Swin Transformer-Tinyself.swin = create_model('swin_tiny_patch4_window7_224', pretrained=True)# 移除最后的分类头self.swin_features = nn.Sequential(*list(self.swin.children())[:-1])# 获取两种模型的输出特征维度self.resnet_out_dim = 2048  # ResNet50的输出维度self.swin_out_dim = 768     # Swin Transformer-Tiny的输出维度# 注意力机制用于特征融合self.attention = nn.Sequential(nn.Linear(self.resnet_out_dim + self.swin_out_dim, 512),nn.ReLU(),nn.Linear(512, 2),nn.Softmax(dim=1))# 最终分类器self.classifier = nn.Sequential(nn.Linear(self.resnet_out_dim + self.swin_out_dim, 512),nn.BatchNorm1d(512),nn.ReLU(),nn.Dropout(0.5),nn.Linear(512, 256),nn.BatchNorm1d(256),nn.ReLU(),nn.Dropout(0.5),nn.Linear(256, num_classes))# 冻结部分预训练层，只微调高层for param in list(self.resnet.parameters())[:-100]:param.requires_grad = Falsefor param in list(self.swin.parameters())[:-100]:param.requires_grad = Falsedef forward(self, x):# ResNet特征提取resnet_feat = self.resnet_features(x)resnet_feat = resnet_feat.view(resnet_feat.size(0), -1)  # 展平# Swin Transformer特征提取swin_feat = self.swin_features(x)swin_feat = swin_feat.view(swin_feat.size(0), -1)  # 展平# 特征融合combined = torch.cat((resnet_feat, swin_feat), dim=1)# 应用注意力机制attn_weights = self.attention(combined)attn_resnet = attn_weights[:, 0].unsqueeze(1) * resnet_featattn_swin = attn_weights[:, 1].unsqueeze(1) * swin_featattn_combined = torch.cat((attn_resnet, attn_swin), dim=1)# 分类out = self.classifier(attn_combined)return out# 训练函数
def train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device):best_val_acc = 0.0train_losses = []val_losses = []train_accs = []val_accs = []for epoch in range(num_epochs):print(f'Epoch {epoch+1}/{num_epochs}')print('-' * 10)# 训练阶段model.train()running_loss = 0.0running_corrects = 0all_preds = []all_labels = []for inputs, labels in train_loader:inputs = inputs.to(device)labels = labels.to(device)# 清零梯度optimizer.zero_grad()# 前向传播outputs = model(inputs)_, preds = torch.max(outputs, 1)loss = criterion(outputs, labels)# 反向传播和优化loss.backward()optimizer.step()# 统计running_loss += loss.item() * inputs.size(0)running_corrects += torch.sum(preds == labels.data)all_preds.extend(preds.cpu().numpy())all_labels.extend(labels.cpu().numpy())# 计算训练集指标epoch_loss = running_loss / len(train_loader.dataset)epoch_acc = running_corrects.double() / len(train_loader.dataset)train_precision = precision_score(all_labels, all_preds)train_recall = recall_score(all_labels, all_preds)train_f1 = f1_score(all_labels, all_preds)train_losses.append(epoch_loss)train_accs.append(epoch_acc.item())print(f'Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f} 'f'Precision: {train_precision:.4f} Recall: {train_recall:.4f} F1: {train_f1:.4f}')# 验证阶段model.eval()val_running_loss = 0.0val_running_corrects = 0val_all_preds = []val_all_labels = []with torch.no_grad():for inputs, labels in val_loader:inputs = inputs.to(device)labels = labels.to(device)outputs = model(inputs)_, preds = torch.max(outputs, 1)loss = criterion(outputs, labels)val_running_loss += loss.item() * inputs.size(0)val_running_corrects += torch.sum(preds == labels.data)val_all_preds.extend(preds.cpu().numpy())val_all_labels.extend(labels.cpu().numpy())# 计算验证集指标val_epoch_loss = val_running_loss / len(val_loader.dataset)val_epoch_acc = val_running_corrects.double() / len(val_loader.dataset)val_precision = precision_score(val_all_labels, val_all_preds)val_recall = recall_score(val_all_labels, val_all_preds)val_f1 = f1_score(val_all_labels, val_all_preds)val_losses.append(val_epoch_loss)val_accs.append(val_epoch_acc.item())print(f'Val Loss: {val_epoch_loss:.4f} Acc: {val_epoch_acc:.4f} 'f'Precision: {val_precision:.4f} Recall: {val_recall:.4f} F1: {val_f1:.4f}')# 学习率调度scheduler.step()# 保存最佳模型if val_epoch_acc > best_val_acc:best_val_acc = val_epoch_acctorch.save(model.state_dict(), 'best_ps_trace_model.pth')print(f'Saved best model with accuracy: {best_val_acc:.4f}')print()# 绘制训练曲线plt.figure(figsize=(12, 4))plt.subplot(1, 2, 1)plt.plot(train_losses, label='Train Loss')plt.plot(val_losses, label='Val Loss')plt.title('Loss Curves')plt.xlabel('Epoch')plt.ylabel('Loss')plt.legend()plt.subplot(1, 2, 2)plt.plot(train_accs, label='Train Accuracy')plt.plot(val_accs, label='Val Accuracy')plt.title('Accuracy Curves')plt.xlabel('Epoch')plt.ylabel('Accuracy')plt.legend()plt.tight_layout()plt.savefig('training_curves.png')plt.close()print(f'Training complete. Best val Acc: {best_val_acc:.4f}')return model# 主函数
def main():# 配置参数data_dir = './ps_dataset'  # 数据集目录batch_size = 16learning_rate = 1e-4num_epochs = 20num_workers = 4# 设备配置device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")print(f'Using device: {device}')# 数据加载train_transform, val_transform = get_transforms()train_dataset = PSTraceDataset(root_dir=data_dir, transform=train_transform, train=True)val_dataset = PSTraceDataset(root_dir=data_dir, transform=val_transform, train=False)train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)print(f'Train samples: {len(train_dataset)}, Val samples: {len(val_dataset)}')# 初始化模型model = PSTraceDetector(num_classes=2)model = model.to(device)# 定义损失函数和优化器criterion = nn.CrossEntropyLoss()optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-5)# 学习率调度器scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)# 训练模型model = train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device)# 加载最佳模型并在验证集上进行最终评估model.load_state_dict(torch.load('best_ps_trace_model.pth'))model.eval()final_preds = []final_labels = []with torch.no_grad():for inputs, labels in val_loader:inputs = inputs.to(device)labels = labels.to(device)outputs = model(inputs)_, preds = torch.max(outputs, 1)final_preds.extend(preds.cpu().numpy())final_labels.extend(labels.cpu().numpy())# 计算最终指标final_acc = accuracy_score(final_labels, final_preds)final_precision = precision_score(final_labels, final_preds)final_recall = recall_score(final_labels, final_preds)final_f1 = f1_score(final_labels, final_preds)print('Final Evaluation on Validation Set:')print(f'Accuracy: {final_acc:.4f}')print(f'Precision: {final_precision:.4f}')print(f'Recall: {final_recall:.4f}')print(f'F1 Score: {final_f1:.4f}')if __name__ == '__main__':main()

代码讲解

1. 数据集处理

代码中定义了PSTraceDataset类来处理PS痕迹识别的数据集，假设数据集结构如下：

ps_dataset/
├── original/    # 原始图像
└── tampered/    # 经过PS处理的图像

数据集类会自动平衡两类图像的数量，并按8:2的比例划分训练集和验证集。

2. 数据增强

为了提高模型的泛化能力，使用了多种数据增强技术：

随机翻转（水平和垂直）
随机旋转
随机仿射变换
颜色抖动（亮度、对比度、饱和度）

3. 模型架构

PSTraceDetector类实现了Swin Transformer-T与ResNet50的融合模型：

特征提取：
- ResNet50提取局部纹理特征，输出维度为2048
- Swin Transformer-T提取全局上下文特征，输出维度为768
特征融合：
- 使用注意力机制自动学习两种特征的权重
- 将加权后的特征拼接，形成融合特征
分类器：
- 采用多层感知机(MLP)进行最终分类
- 加入批归一化和dropout层防止过拟合
迁移学习：
- 使用预训练权重初始化模型
- 冻结部分底层参数，只微调高层参数

4. 训练过程

训练函数实现了完整的训练和验证流程：

使用交叉熵损失函数
采用AdamW优化器和学习率调度
跟踪多种评估指标（准确率、精确率、召回率、F1分数）
保存性能最佳的模型
绘制训练曲线（损失和准确率）

模型分析

优势分析

混合架构优势：
- ResNet50擅长捕捉图像的局部特征和边缘信息，对检测细微的PS痕迹至关重要
- Swin Transformer能够建模长距离依赖关系，有助于发现图像中不一致的区域
- 两者结合可以弥补单一模型的不足
注意力融合机制：
- 动态调整两种特征的权重，在不同场景下自动侧重更重要的特征源
- 提高模型对复杂PS操作的识别能力
迁移学习策略：
- 利用预训练模型的特征提取能力，加速收敛并提高性能
- 选择性冻结底层参数，避免过拟合并减少计算量