资讯详情

点击查看代码

import numpy as np
import matplotlib.pyplot as plt
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms# ======================== 1. 设备配置与参数设置 ========================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size = 128
num_classes = 10  # 模型输出维度：0-9（对应数字1-10）
num_epochs = 5
lr = 0.001# ======================== 2. 数据预处理 ========================
# 数据变换：张量转换 + 标准化
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,))  # MNIST官方均值和标准差
])# 加载MNIST数据集（原始标签0-9，对应数字1-10）
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)# 注意：不再偏移标签！训练时用原始0-9标签（符合CrossEntropyLoss要求），仅在评估/展示时映射为1-10
# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)# ======================== 3. 模型定义 ========================
# 3.1 共享CNN特征提取器
class CNNFeatureExtractor(nn.Module):def __init__(self):super(CNNFeatureExtractor, self).__init__()self.features = nn.Sequential(nn.Conv2d(1, 32, kernel_size=3, padding=1),nn.ReLU(),nn.MaxPool2d(2, 2),  # 输出: 32 * 14 * 14nn.Conv2d(32, 64, kernel_size=3, padding=1),nn.ReLU(),nn.MaxPool2d(2, 2),  # 输出: 64 * 7 * 7nn.Conv2d(64, 128, kernel_size=3, padding=1),nn.ReLU(),nn.MaxPool2d(2, 2)  # 输出: 128 * 3 * 3)self.flatten_dim = 128 * 3 * 3  # 展平后维度def forward(self, x):x = self.features(x)x = x.view(-1, self.flatten_dim)return x# 3.2 CNN+Softmax
class CNNSoftmax(nn.Module):def __init__(self, num_classes):super(CNNSoftmax, self).__init__()self.feature_extractor = CNNFeatureExtractor()self.classifier = nn.Linear(self.feature_extractor.flatten_dim, num_classes)self.softmax = nn.Softmax(dim=1)def forward(self, x):features = self.feature_extractor(x)logits = self.classifier(features)return self.softmax(logits), logits# 3.3 CNN+Sigmoid
class CNNSigmoid(nn.Module):def __init__(self, num_classes):super(CNNSigmoid, self).__init__()self.feature_extractor = CNNFeatureExtractor()self.classifier = nn.Linear(self.feature_extractor.flatten_dim, num_classes)self.sigmoid = nn.Sigmoid()def forward(self, x):features = self.feature_extractor(x)logits = self.classifier(features)return self.sigmoid(logits), logits# 3.4 CNN+SVM（多分类Hinge Loss）
class CNNSVM(nn.Module):def __init__(self, num_classes):super(CNNSVM, self).__init__()self.feature_extractor = CNNFeatureExtractor()self.classifier = nn.Linear(self.feature_extractor.flatten_dim, num_classes)def forward(self, x):features = self.feature_extractor(x)return self.classifier(features)# 多分类SVM Hinge Loss定义
class MultiClassHingeLoss(nn.Module):def __init__(self, margin=1.0):super(MultiClassHingeLoss, self).__init__()self.margin = margindef forward(self, logits, labels):# labels已经是0-9，无需调整one_hot = torch.zeros_like(logits).scatter(1, labels.unsqueeze(1), 1)correct_logit = torch.sum(logits * one_hot, dim=1, keepdim=True)loss = torch.maximum(torch.tensor(0.0).to(device), self.margin - (correct_logit - logits))loss = torch.sum(loss * (1 - one_hot)) / logits.size(0)return loss# ======================== 4. 训练与评估函数 ========================
# 4.1 通用CNN训练函数
def train_cnn_model(model, criterion, optimizer, train_loader, num_epochs, is_sigmoid=False):model.to(device)model.train()start_time = time.time()for epoch in range(num_epochs):total_loss = 0.0for batch_idx, (data, target) in enumerate(train_loader):data, target = data.to(device), target.to(device)  # target是0-9，符合要求optimizer.zero_grad()if is_sigmoid:# Sigmoid需要one-hot标签（0-9）target_onehot = torch.zeros(len(target), num_classes).to(device).scatter(1, target.unsqueeze(1), 1)output, logits = model(data)loss = criterion(logits, target_onehot)else:if isinstance(model, (CNNSoftmax, CNNSigmoid)):output, logits = model(data)else:logits = model(data)loss = criterion(logits, target)loss.backward()optimizer.step()total_loss += loss.item() * data.size(0)avg_loss = total_loss / len(train_loader.dataset)print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss:.4f}')train_time = time.time() - start_timereturn model, train_time# 4.2 模型评估函数（计算准确率、精确率、召回率、F1）
def evaluate_model(model, test_loader, is_cnn_svm=False):model.eval()all_preds = []all_targets = []with torch.no_grad():for data, target in test_loader:data, target = data.to(device), target.to(device)  # target是0-9if isinstance(model, CNNSoftmax):output, _ = model(data)preds = torch.argmax(output, dim=1)  # 0-9elif isinstance(model, CNNSigmoid):output, _ = model(data)preds = torch.argmax(output, dim=1)  # 0-9elif is_cnn_svm:logits = model(data)preds = torch.argmax(logits, dim=1)  # 0-9# 评估时映射为1-10（仅为了符合题目“1-10识别”的表述）all_preds.extend((preds + 1).cpu().numpy())  # 0-9 → 1-10all_targets.extend((target + 1).cpu().numpy())  # 0-9 → 1-10all_preds = np.array(all_preds)all_targets = np.array(all_targets)# 手动实现宏观平均的精确率、召回率、F1（避免依赖sklearn）def calculate_metrics(y_true, y_pred, num_classes):# 转换回0-9计算指标（避免10超出索引）y_true = y_true - 1y_pred = y_pred - 1# 初始化各类的TP、FP、FNTP = np.zeros(num_classes)FP = np.zeros(num_classes)FN = np.zeros(num_classes)for cls in range(num_classes):TP[cls] = np.sum((y_true == cls) & (y_pred == cls))FP[cls] = np.sum((y_true != cls) & (y_pred == cls))FN[cls] = np.sum((y_true == cls) & (y_pred != cls))# 计算各类的精确率、召回率（避免除以0）precision_per_cls = TP / (TP + FP + 1e-8)recall_per_cls = TP / (TP + FN + 1e-8)f1_per_cls = 2 * (precision_per_cls * recall_per_cls) / (precision_per_cls + recall_per_cls + 1e-8)# 宏观平均precision = np.mean(precision_per_cls)recall = np.mean(recall_per_cls)f1 = np.mean(f1_per_cls)accuracy = np.sum(TP) / len(y_true)return accuracy, precision, recall, f1accuracy, precision, recall, f1 = calculate_metrics(all_targets, all_preds, num_classes)return accuracy, precision, recall, f1# ======================== 5. 模型训练 ========================
print("=" * 30 + " Training CNN+Softmax " + "=" * 30)
model_softmax = CNNSoftmax(num_classes)
criterion_softmax = nn.CrossEntropyLoss()
optimizer_softmax = optim.Adam(model_softmax.parameters(), lr=lr)
model_softmax, time_softmax = train_cnn_model(model_softmax, criterion_softmax, optimizer_softmax, train_loader,num_epochs)print("\n" + "=" * 30 + " Training CNN+Sigmoid " + "=" * 30)
model_sigmoid = CNNSigmoid(num_classes)
criterion_sigmoid = nn.BCEWithLogitsLoss()
optimizer_sigmoid = optim.Adam(model_sigmoid.parameters(), lr=lr)
model_sigmoid, time_sigmoid = train_cnn_model(model_sigmoid, criterion_sigmoid, optimizer_sigmoid, train_loader,num_epochs, is_sigmoid=True)print("\n" + "=" * 30 + " Training CNN+SVM " + "=" * 30)
model_cnn_svm = CNNSVM(num_classes)
criterion_cnn_svm = MultiClassHingeLoss()
optimizer_cnn_svm = optim.Adam(model_cnn_svm.parameters(), lr=lr)
model_cnn_svm, time_cnn_svm = train_cnn_model(model_cnn_svm, criterion_cnn_svm, optimizer_cnn_svm, train_loader,num_epochs)# ======================== 6. 模型评估 ========================
print("\n" + "=" * 30 + " Model Evaluation " + "=" * 30)
# 评估三个CNN模型
acc_softmax, pre_softmax, rec_softmax, f1_softmax = evaluate_model(model_softmax, test_loader)
acc_sigmoid, pre_sigmoid, rec_sigmoid, f1_sigmoid = evaluate_model(model_sigmoid, test_loader)
acc_cnn_svm, pre_cnn_svm, rec_cnn_svm, f1_cnn_svm = evaluate_model(model_cnn_svm, test_loader, is_cnn_svm=True)# 整理结果
models = ['CNN+Softmax', 'CNN+Sigmoid', 'CNN+SVM']
accuracy_list = [acc_softmax, acc_sigmoid, acc_cnn_svm]
precision_list = [pre_softmax, pre_sigmoid, pre_cnn_svm]
recall_list = [rec_softmax, rec_sigmoid, rec_cnn_svm]
f1_list = [f1_softmax, f1_sigmoid, f1_cnn_svm]
time_list = [time_softmax, time_sigmoid, time_cnn_svm]# 打印结果表格
print("\n" + "=" * 90)
print(f"{'Model':<18} {'Accuracy':<12} {'Precision':<12} {'Recall':<12} {'F1-Score':<12} {'Train Time(s)':<12}")
print("=" * 90)
for i in range(len(models)):print(f"{models[i]:<18} {accuracy_list[i]:<12.4f} {precision_list[i]:<12.4f} {recall_list[i]:<12.4f} {f1_list[i]:<12.4f} {time_list[i]:<12.2f}")
print("=" * 90)# ======================== 7. 结果可视化 ========================
def plot_metrics_comparison(models, accuracy, precision, recall, f1, time_list):fig, ax = plt.subplots(1, 2, figsize=(18, 7))x = np.arange(len(models))width = 0.2# 子图1：分类指标对比ax[0].bar(x - 1.5 * width, accuracy, width, label='Accuracy', color='#1f77b4', edgecolor='black')ax[0].bar(x - 0.5 * width, precision, width, label='Precision', color='#ff7f0e', edgecolor='black')ax[0].bar(x + 0.5 * width, recall, width, label='Recall', color='#2ca02c', edgecolor='black')ax[0].bar(x + 1.5 * width, f1, width, label='F1-Score', color='#d62728', edgecolor='black')ax[0].set_xlabel('Models', fontsize=12, fontweight='bold')ax[0].set_ylabel('Score', fontsize=12, fontweight='bold')ax[0].set_title('Classification Metrics Comparison (Digits 1-10)', fontsize=14, fontweight='bold')ax[0].set_xticks(x)ax[0].set_xticklabels(models, rotation=15, fontsize=10)ax[0].legend(fontsize=10)ax[0].grid(axis='y', linestyle='--', alpha=0.7)# 子图2：训练时间对比bars = ax[1].bar(models, time_list, color=['#1f77b4', '#ff7f0e', '#2ca02c'], edgecolor='black')ax[1].set_xlabel('Models', fontsize=12, fontweight='bold')ax[1].set_ylabel('Training Time (s)', fontsize=12, fontweight='bold')ax[1].set_title('Training Time Comparison', fontsize=14, fontweight='bold')ax[1].tick_params(axis='x', rotation=15, labelsize=10)ax[1].grid(axis='y', linestyle='--', alpha=0.7)# 添加数值标签for bar in bars:height = bar.get_height()ax[1].text(bar.get_x() + bar.get_width() / 2., height + 0.5, f'{height:.2f}', ha='center', va='bottom',fontsize=10)plt.tight_layout()plt.savefig('mnist_digit_recognition_comparison.png', dpi=300, bbox_inches='tight')plt.show()# 调用绘图函数
plot_metrics_comparison(models, accuracy_list, precision_list, recall_list, f1_list, time_list)
print("2024310143014")