1. PyTorch模型参数量与计算量统计

任务目标

训练一个简单的CNN模型完成MNIST手写数字分类，保存模型权重并记录其测试准确率与模型文件大小，作为后续量化压缩实验的基准参照。

环境依赖

• Python 3.8+
• PyTorch 1.13+（含torchvision）
• numpy

完整实现代码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import os
import time

# 1. 定义CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 64 * 7 * 7)
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 2. 数据加载
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, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)

# 3. 训练配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 4. 训练函数
def train(epochs=3):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f'Epoch {epoch+1}/{epochs}, Loss: {running_loss/len(train_loader):.4f}')

# 5. 测试函数
def evaluate():
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = 100.0 * correct / total
    return accuracy

# 6. 执行训练与评估
if __name__ == '__main__':
    start_time = time.time()
    train(epochs=3)
    accuracy = evaluate()
    train_time = time.time() - start_time

    # 保存模型权重
    torch.save(model.state_dict(), 'mnist_cnn.pth')
    model_size = os.path.getsize('mnist_cnn.pth') / 1024.0  # KB

    print(f'测试准确率: {accuracy:.2f}%')
    print(f'模型文件大小: {model_size:.2f} KB')
    print(f'训练耗时: {train_time:.2f} 秒')

运行方式

python mnist_train.py

验证方法与判定标准

• 运行成功标志：程序正常结束，控制台输出每个epoch的训练损失，并最终打印测试准确率、模型文件大小和训练耗时。
• 基准记录：记录输出的测试准确率（预期约98%～99%）和模型文件大小（约430KB），作为后续量化实验的对比基准。
• 文件检查：当前目录下生成mnist_cnn.pth文件，可通过ls -lh mnist_cnn.pth确认文件存在且大小与输出一致。

关键说明

• 首次运行会自动下载MNIST数据集至./data目录，需保持网络畅通。
• 模型结构为两层卷积+两层全连接，参数量约43万，属于轻量级CNN，适合作为量化基准模型。
• 训练轮数设为3轮以控制时间，若需更高精度可增加至5轮，但基准记录需保持一致。