import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms import numpy as np import time import matplotlib.pyplot as plt # 设置随机种子确保结果可复现 torch.manual_seed(42) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # 加载MNIST数据集 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform) test_dataset = datasets.MNIST('./data', train=False, transform=transform) train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1000) # 定义一个简单的CNN模型 class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) self.fc1 = nn.Linear(64 * 7 * 7, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2(x), 2)) x = x.view(-1, 64 * 7 * 7) x = F.relu(self.fc1(x)) x = self.fc2(x) return F.log_softmax(x, dim=1) # 计算模型参数量 def count_parameters(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) # 训练函数 def train(model, train_loader, optimizer, epoch): model.train() train_loss = 0 correct = 0 for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.nll_loss(output, target) loss.backward() optimizer.step() train_loss += loss.item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() train_loss /= len(train_loader.dataset) accuracy = 100. * correct / len(train_loader.dataset) print(f'Train Epoch: {epoch} Loss: {train_loss:.4f} Accuracy: {accuracy:.2f}%') return accuracy # 测试函数 def test(model, test_loader): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += F.nll_loss(output, target, reduction='sum').item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) accuracy = 100. * correct / len(test_loader.dataset) print(f'Test Loss: {test_loss:.4f} Accuracy: {accuracy:.2f}%') return accuracy # 评估推理速度 def evaluate_inference_time(model, test_loader): model.eval() total_time = 0 num_batches = 0 with torch.no_grad(): for data, _ in test_loader: data = data.to(device) start_time = time.time() _ = model(data) end_time = time.time() total_time += (end_time - start_time) num_batches += 1 avg_time = total_time / num_batches print(f"Average inference time per batch: {avg_time*1000:.2f} ms") return avg_time # 结构化剪枝函数 (通道剪枝) def channel_prune(model, prune_ratio=0.5): # 对conv1层进行剪枝 conv1_weight = model.conv1.weight.data.clone() # 计算每个通道的L1范数 l1_norm = torch.sum(torch.abs(conv1_weight), dim=[1, 2, 3]) # 确定阈值，保留(1-prune_ratio)的通道 num_channels = conv1_weight.size(0) num_to_keep = int(num_channels * (1 - prune_ratio)) # 排序找出需要保留的通道索引 _, indices = torch.sort(l1_norm) keep_indices = indices[-num_to_keep:] # 创建新模型，只保留选定的通道 pruned_model = SimpleCNN().to(device) # 重新初始化第一个卷积层，只保留选定的通道 pruned_model.conv1 = nn.Conv2d(1, num_to_keep, kernel_size=3, stride=1, padding=1).to(device) pruned_model.conv1.weight.data = torch.index_select(model.conv1.weight.data, 0, keep_indices) if model.conv1.bias is not None: pruned_model.conv1.bias.data = torch.index_select(model.conv1.bias.data, 0, keep_indices) # 重新初始化第二个卷积层，考虑到输入通道的变化 pruned_model.conv2 = nn.Conv2d(num_to_keep, 64, kernel_size=3, stride=1, padding=1).to(device) pruned_model.conv2.weight.data = torch.randn_like(pruned_model.conv2.weight.data) * 0.01 # 返回剪枝后的模型 return pruned_model # 训练原始模型 print("Training original model...") model = SimpleCNN().to(device) optimizer = optim.Adam(model.parameters(), lr=0.001) original_train_accuracies = [] original_test_accuracies = [] for epoch in range(1, 6): train_acc = train(model, train_loader, optimizer, epoch) test_acc = test(model, test_loader) original_train_accuracies.append(train_acc) original_test_accuracies.append(test_acc) original_params = count_parameters(model) print(f"Original model parameters: {original_params}") # 评估原始模型的推理时间 original_inference_time = evaluate_inference_time(model, test_loader) # 进行结构化剪枝 print("\nPruning model...") pruned_model = channel_prune(model, prune_ratio=0.5) # 评估剪枝后模型的初始性能 print("\nEvaluating pruned model before fine-tuning...") pruned_test_acc_before = test(pruned_model, test_loader) # 微调剪枝后的模型 print("\nFine-tuning pruned model...") pruned_optimizer = optim.Adam(pruned_model.parameters(), lr=0.001) pruned_train_accuracies = [] pruned_test_accuracies = [] for epoch in range(1, 11): train_acc = train(pruned_model, train_loader, pruned_optimizer, epoch) test_acc = test(pruned_model, test_loader) pruned_train_accuracies.append(train_acc) pruned_test_accuracies.append(test_acc) # 计算剪枝后的参数量 pruned_params = count_parameters(pruned_model) print(f"Pruned model parameters: {pruned_params}") print(f"Parameter reduction: {100 * (1 - pruned_params/original_params):.2f}%") # 评估剪枝后模型的推理时间 pruned_inference_time = evaluate_inference_time(pruned_model, test_loader) print(f"Inference time reduction: {100 * (1 - pruned_inference_time/original_inference_time):.2f}%") # 可视化原始和剪枝后模型的准确率 plt.figure(figsize=(10, 6)) plt.plot(range(1, 6), original_test_accuracies, 'b-', label='Original Model') plt.plot(range(1, 11), pruned_test_accuracies, 'r-', label='Pruned Model') plt.axhline(y=original_test_accuracies[-1], color='b', linestyle='--', label='Original Final Accuracy') plt.xlabel('Epochs') plt.ylabel('Test Accuracy (%)') plt.title('Test Accuracy Comparison') plt.legend() plt.grid(True) plt.savefig('pruning_comparison.png') plt.close() print("\nResults Summary:") print(f"Original Model - Parameters: {original_params}, Test Accuracy: {original_test_accuracies[-1]:.2f}%, Inference Time: {original_inference_time*1000:.2f} ms") print(f"Pruned Model - Parameters: {pruned_params}, Test Accuracy: {pruned_test_accuracies[-1]:.2f}%, Inference Time: {pruned_inference_time*1000:.2f} ms")