CIFAR-10图像分类

1
import torch
2
from torchvision import transforms, datasets
3
from torch.utils.data import DataLoader, random_split
4

5

6
class CIFAR10Data:
7

8
    def __init__(self, data_dir='./data', batch_size=64, num_workers=2):
9
        self.data_dir = data_dir
10
        self.batch_size = batch_size
11
        self.num_workers = num_workers
12

13
        # CIFAR-10 预计算的通道均值与标准差，用于标准化
14
        self.mean = (0.4914, 0.4822, 0.4465)
15
        self.std = (0.2023, 0.1994, 0.2010)
16

17
        # 训练集 transform：随机裁剪 + 水平翻转做数据增强，提升泛化能力
18
        self.train_transform = transforms.Compose([
19
            transforms.RandomCrop(32, padding=4),
20
            transforms.RandomHorizontalFlip(p=0.5),
21
            transforms.ToTensor(),
22
            transforms.Normalize(self.mean, self.std)
23
        ])
24

25
        # 测试集 transform：只做 ToTensor 与归一化，不引入随机增强
26
        self.test_transform = transforms.Compose([
27
            transforms.ToTensor(),
28
            transforms.Normalize(self.mean, self.std)
29
        ])
30

31
    def load_data(self):
32
        # 分别加载训练集（50000 张）与测试集（10000 张），本地缺失时自动下载
33
        full_train = datasets.CIFAR10(
34
            root=self.data_dir,
35
            train=True,
36
            download=True,
37
            transform=self.train_transform
38
        )
39

40
        test_dataset = datasets.CIFAR10(
41
            root=self.data_dir,
42
            train=False,
43
            download=True,
44
            transform=self.test_transform
45
        )
46

47
        print(f"训练集大小: {len(full_train)}")
48
        print(f"测试集大小: {len(test_dataset)}")
49

50
        return full_train, test_dataset
51

52
    def split_data(self, full_train, val_ratio=0.1):
53
        # 按 9:1 划分训练/验证子集，固定随机种子保证可复现
54
        total_size = len(full_train)
55
        val_size = int(total_size * val_ratio)
56
        train_size = total_size - val_size
57

58
        train_dataset, val_dataset = random_split(
59
            full_train,
60
            [train_size, val_size],
61
            generator=torch.Generator().manual_seed(42)
62
        )
63

64
        # random_split 返回 Subset 且不复制数据，需改回原数据集的 transform 去掉增强
65
        val_dataset.dataset.transform = self.test_transform
66

67
        print(f"训练子集大小: {len(train_dataset)}")
68
        print(f"验证子集大小: {len(val_dataset)}")
69

70
        return train_dataset, val_dataset
71

72
    def get_loaders(self, val_ratio=0.1):
73
        full_train, test_dataset = self.load_data()
74
        train_dataset, val_dataset = self.split_data(full_train, val_ratio)
75

76
        # 训练集打乱顺序，验证/测试集保持顺序；pin_memory 在 CUDA 可用时开启以加速传输
77
        train_loader = DataLoader(
78
            train_dataset,
79
            batch_size=self.batch_size,
80
            shuffle=True,
81
            num_workers=self.num_workers,
82
            pin_memory=True if torch.cuda.is_available() else False
83
        )
84

85
        val_loader = DataLoader(
86
            val_dataset,
87
            batch_size=self.batch_size,
88
            shuffle=False,
89
            num_workers=self.num_workers,
90
            pin_memory=True if torch.cuda.is_available() else False
91
        )
92

93
        test_loader = DataLoader(
94
            test_dataset,
95
            batch_size=self.batch_size,
96
            shuffle=False,
97
            num_workers=self.num_workers,
98
            pin_memory=True if torch.cuda.is_available() else False
99
        )
100

101
        return train_loader, val_loader, test_loader
102

103

104
if __name__ == '__main__':
105
    data_manager = CIFAR10Data(batch_size=64)
106
    train_loader, val_loader, test_loader = data_manager.get_loaders()
107

108
    # 取出一个 batch 检查数据管道
109
    images, labels = next(iter(train_loader))
110
    print(f"图像 batch 形状: {images.shape}")
111
    print(f"标签 batch 形状: {labels.shape}")
112
    print(f"图像数值范围: [{images.min():.2f}, {images.max():.2f}]")

1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4

5

6
class MLP(nn.Module):
7

8
    def __init__(self, num_classes=10, dropout=0.3):
9
        super(MLP, self).__init__()
10

11
        self.flatten = nn.Flatten()
12
        self.fc1 = nn.Linear(32 * 32 * 3, 512)
13
        self.dropout1 = nn.Dropout(dropout)
14
        self.fc2 = nn.Linear(512, 256)
15
        self.dropout2 = nn.Dropout(dropout)
16
        self.fc3 = nn.Linear(256, num_classes)
17

18
        self._initialize_weights()
19

20
    def _initialize_weights(self):
21
        # Kaiming 初始化适配 ReLU 激活，偏置置零
22
        for m in self.modules():
23
            if isinstance(m, nn.Linear):
24
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
25
                if m.bias is not None:
26
                    nn.init.constant_(m.bias, 0)
27

28
    def forward(self, x):
29
        x = self.flatten(x)  # [batch, 3, 32, 32] 展平为 [batch, 3072]
30
        x = F.relu(self.fc1(x))
31
        x = self.dropout1(x)
32
        x = F.relu(self.fc2(x))
33
        x = self.dropout2(x)
34
        x = self.fc3(x)  # 只输出 logits，不加 Softmax，交给 CrossEntropyLoss 处理
35
        return x
36

37

38
class BasicBlock(nn.Module):
39

40
    def __init__(self, in_channels, out_channels, stride=1):
41
        super(BasicBlock, self).__init__()
42

43
        self.conv1 = nn.Conv2d(
44
            in_channels, out_channels, kernel_size=3,
45
            stride=stride, padding=1, bias=False
46
        )
47
        self.bn1 = nn.BatchNorm2d(out_channels)
48

49
        self.conv2 = nn.Conv2d(
50
            out_channels, out_channels, kernel_size=3,
51
            stride=1, padding=1, bias=False
52
        )
53
        self.bn2 = nn.BatchNorm2d(out_channels)
54

55
        # 输入输出维度不一致时，用 1x1 卷积构建投影 shortcut 对齐维度
56
        self.shortcut = nn.Sequential()
57
        if stride != 1 or in_channels != out_channels:
58
            self.shortcut = nn.Sequential(
59
                nn.Conv2d(in_channels, out_channels, kernel_size=1,
60
                          stride=stride, bias=False),
61
                nn.BatchNorm2d(out_channels)
62
            )
63

64
    def forward(self, x):
65
        identity = self.shortcut(x)
66
        out = self.conv1(x)
67
        out = self.bn1(out)
68
        out = F.relu(out)
69
        out = self.conv2(out)
70
        out = self.bn2(out)
71
        out += identity  # 主干输出与 shortcut 相加，完成残差连接
72
        out = F.relu(out)
73
        return out
74

75

76
class ResNet18(nn.Module):
77

78
    def __init__(self, num_classes=10):
79
        super(ResNet18, self).__init__()
80

81
        # 针对 32x32 小图像，用 3x3 步长 1 卷积替代原版 7x7 步长 2 卷积
82
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
83
        self.bn1 = nn.BatchNorm2d(64)
84

85
        # 4 个残差层通道数 64 -> 128 -> 256 -> 512，stride=2 时空间尺寸减半
86
        self.layer1 = self._make_layer(64, 64, num_blocks=2, stride=1)
87
        self.layer2 = self._make_layer(64, 128, num_blocks=2, stride=2)
88
        self.layer3 = self._make_layer(128, 256, num_blocks=2, stride=2)
89
        self.layer4 = self._make_layer(256, 512, num_blocks=2, stride=2)
90

91
        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
92
        self.fc = nn.Linear(512, num_classes)
93

94
        self._initialize_weights()
95

96
    def _make_layer(self, in_channels, out_channels, num_blocks, stride):
97
        layers = []
98
        # 第一个块负责下采样，后续块保持特征图尺寸不变
99
        layers.append(BasicBlock(in_channels, out_channels, stride))
100
        for _ in range(1, num_blocks):
101
            layers.append(BasicBlock(out_channels, out_channels, stride=1))
102
        return nn.Sequential(*layers)
103

104
    def _initialize_weights(self):
105
        for m in self.modules():
106
            if isinstance(m, nn.Conv2d):
107
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
108
            elif isinstance(m, nn.BatchNorm2d):
109
                nn.init.constant_(m.weight, 1)
110
                nn.init.constant_(m.bias, 0)
111

112
    def forward(self, x):
113
        x = F.relu(self.bn1(self.conv1(x)))
114
        x = self.layer1(x)
115
        x = self.layer2(x)
116
        x = self.layer3(x)
117
        x = self.layer4(x)
118
        x = self.avg_pool(x)
119
        x = x.view(x.size(0), -1)  # 全局池化后展平为 [batch, 512]
120
        x = self.fc(x)
121
        return x
122

123

124
if __name__ == '__main__':
125
    mlp = MLP(num_classes=10)
126
    x = torch.randn(2, 3, 32, 32)
127
    out = mlp(x)
128
    print(f"MLP 输出形状: {out.shape}")
129
    print(f"MLP 参数量: {sum(p.numel() for p in mlp.parameters()) / 1e6:.2f}M")
130

131
    resnet = ResNet18(num_classes=10)
132
    out = resnet(x)
133
    print(f"ResNet 输出形状: {out.shape}")
134
    print(f"ResNet 参数量: {sum(p.numel() for p in resnet.parameters()) / 1e6:.2f}M")

1
import torch
2
import torch.nn as nn
3
import torch.optim as optim
4
from tqdm import tqdm
5
import json
6
import os
7
import matplotlib.pyplot as plt
8
from sklearn.metrics import confusion_matrix, classification_report
9
import numpy as np
10

11

12
class Trainer:
13

14
    def __init__(self, model, train_loader, val_loader, device,
15
                 lr=0.001, weight_decay=1e-4, model_name='model'):
16
        self.model = model.to(device)
17
        self.train_loader = train_loader
18
        self.val_loader = val_loader
19
        self.device = device
20
        self.model_name = model_name
21

22
        # 交叉熵损失 + Adam 优化器（带权重衰减）
23
        self.criterion = nn.CrossEntropyLoss()
24
        self.optimizer = optim.Adam(self.model.parameters(), lr=lr, weight_decay=weight_decay)
25

26
        # 验证准确率连续 3 轮不提升时学习率减半（PyTorch 2.7.1 无 verbose 参数）
27
        self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
28
            self.optimizer,
29
            mode='max',
30
            factor=0.5,
31
            patience=3
32
        )
33

34
        # 训练历史，用于绘制训练曲线
35
        self.history = {
36
            'train_loss': [], 'train_acc': [],
37
            'val_loss': [], 'val_acc': []
38
        }
39
        self.best_acc = 0.0
40
        self.current_lr = lr
41

42
    def train_epoch(self):
43
        self.model.train()
44
        total_loss = 0.0
45
        correct = 0
46
        total = 0
47

48
        pbar = tqdm(self.train_loader, desc='Training')
49
        for inputs, labels in pbar:
50
            inputs, labels = inputs.to(self.device), labels.to(self.device)
51

52
            # 训练五步：梯度清零 -> 前向传播 -> 计算损失 -> 反向传播 -> 参数更新
53
            self.optimizer.zero_grad()
54
            outputs = self.model(inputs)
55
            loss = self.criterion(outputs, labels)
56
            loss.backward()
57
            self.optimizer.step()
58

59
            # 统计当前 epoch 的损失与准确率
60
            total_loss += loss.item()
61
            _, predicted = outputs.max(1)
62
            total += labels.size(0)
63
            correct += predicted.eq(labels).sum().item()
64

65
            pbar.set_postfix({
66
                'loss': f'{loss.item():.4f}',
67
                'acc': f'{100. * correct / total:.2f}%',
68
                'lr': f'{self.optimizer.param_groups[0]["lr"]:.6f}'
69
            })
70

71
        avg_loss = total_loss / len(self.train_loader)
72
        accuracy = 100. * correct / total
73
        return avg_loss, accuracy
74

75
    def validate(self):
76
        # eval 模式 + no_grad：关闭 Dropout 与 BN 更新，不计算梯度
77
        self.model.eval()
78
        total_loss = 0.0
79
        correct = 0
80
        total = 0
81

82
        with torch.no_grad():
83
            pbar = tqdm(self.val_loader, desc='Validating')
84
            for inputs, labels in pbar:
85
                inputs, labels = inputs.to(self.device), labels.to(self.device)
86

87
                outputs = self.model(inputs)
88
                loss = self.criterion(outputs, labels)
89

90
                total_loss += loss.item()
91
                _, predicted = outputs.max(1)
92
                total += labels.size(0)
93
                correct += predicted.eq(labels).sum().item()
94

95
                pbar.set_postfix({'acc': f'{100. * correct / total:.2f}%'})
96

97
        avg_loss = total_loss / len(self.val_loader)
98
        accuracy = 100. * correct / total
99
        return avg_loss, accuracy
100

101
    def train(self, epochs=20, save_dir='./checkpoints'):
102
        os.makedirs(save_dir, exist_ok=True)
103
        best_model_path = os.path.join(save_dir, f'{self.model_name}_best.pth')
104

105
        print(f"\n开始训练 {self.model_name}")
106
        print(f"设备: {self.device}")
107
        print(f"初始学习率: {self.optimizer.param_groups[0]['lr']}")
108
        print(f"总轮数: {epochs}")
109
        print("=" * 50)
110

111
        for epoch in range(epochs):
112
            print(f"\nEpoch [{epoch + 1}/{epochs}]")
113

114
            train_loss, train_acc = self.train_epoch()
115
            val_loss, val_acc = self.validate()
116

117
            self.history['train_loss'].append(train_loss)
118
            self.history['train_acc'].append(train_acc)
119
            self.history['val_loss'].append(val_loss)
120
            self.history['val_acc'].append(val_acc)
121

122
            # 调度器按验证准确率调整学习率，变化时手动打印
123
            old_lr = self.optimizer.param_groups[0]['lr']
124
            self.scheduler.step(val_acc)
125
            new_lr = self.optimizer.param_groups[0]['lr']
126

127
            if new_lr != old_lr:
128
                print(f"! 学习率调整: {old_lr:.6f} -> {new_lr:.6f}")
129

130
            print(f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.2f}%")
131
            print(f"Val   Loss: {val_loss:.4f}, Acc: {val_acc:.2f}%")
132
            print(f"LR: {new_lr:.6f}")
133

134
            # 验证准确率创新高时保存当前权重
135
            if val_acc > self.best_acc:
136
                self.best_acc = val_acc
137
                torch.save(self.model.state_dict(), best_model_path)
138
                print(f"[OK] 保存最佳模型，准确率: {val_acc:.2f}%")
139

140
        # 保存训练历史并绘制损失/准确率曲线
141
        history_path = os.path.join(save_dir, f'{self.model_name}_history.json')
142
        with open(history_path, 'w') as f:
143
            json.dump(self.history, f, indent=2)
144

145
        self.plot_history(save_dir)
146

147
        return self.history
148

149
    def plot_history(self, save_dir):
150
        fig, axes = plt.subplots(1, 2, figsize=(12, 4))
151

152
        axes[0].plot(self.history['train_loss'], label='Train Loss')
153
        axes[0].plot(self.history['val_loss'], label='Val Loss')
154
        axes[0].set_xlabel('Epoch')
155
        axes[0].set_ylabel('Loss')
156
        axes[0].set_title('Loss Curve')
157
        axes[0].legend()
158
        axes[0].grid(True)
159

160
        axes[1].plot(self.history['train_acc'], label='Train Acc')
161
        axes[1].plot(self.history['val_acc'], label='Val Acc')
162
        axes[1].set_xlabel('Epoch')
163
        axes[1].set_ylabel('Accuracy (%)')
164
        axes[1].set_title('Accuracy Curve')
165
        axes[1].legend()
166
        axes[1].grid(True)
167

168
        plt.tight_layout()
169
        plt.savefig(os.path.join(save_dir, f'{self.model_name}_curves.png'), dpi=150)
170
        plt.close()
171

172

173
def test(model, test_loader, device, model_path=None, num_classes=10):
174
    # 加载训练好的最优权重，map_location 保证跨设备加载
175
    if model_path is not None:
176
        model.load_state_dict(torch.load(model_path, map_location=device))
177
        print(f"加载模型: {model_path}")
178

179
    model = model.to(device)
180
    model.eval()
181

182
    all_preds = []
183
    all_labels = []
184
    correct = 0
185
    total = 0
186

187
    # 推理阶段不计算梯度，收集全部预测用于生成分类报告与混淆矩阵
188
    with torch.no_grad():
189
        for inputs, labels in tqdm(test_loader, desc='Testing'):
190
            inputs, labels = inputs.to(device), labels.to(device)
191
            outputs = model(inputs)
192
            _, predicted = outputs.max(1)
193

194
            total += labels.size(0)
195
            correct += predicted.eq(labels).sum().item()
196

197
            all_preds.extend(predicted.cpu().numpy())
198
            all_labels.extend(labels.cpu().numpy())
199

200
    accuracy = 100. * correct / total
201
    print(f"\n测试准确率: {accuracy:.2f}%")
202

203
    classes = ('plane', 'car', 'bird', 'cat', 'deer',
204
               'dog', 'frog', 'horse', 'ship', 'truck')
205
    print("\n分类报告:")
206
    print(classification_report(all_labels, all_preds, target_names=classes, digits=4))
207

208
    # 混淆矩阵热力图：行为真实标签，列为预测标签
209
    cm = confusion_matrix(all_labels, all_preds)
210
    plt.figure(figsize=(10, 8))
211
    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
212
    plt.title('Confusion Matrix')
213
    plt.colorbar()
214
    tick_marks = np.arange(num_classes)
215
    plt.xticks(tick_marks, classes, rotation=45)
216
    plt.yticks(tick_marks, classes)
217
    plt.ylabel('True Label')
218
    plt.xlabel('Predicted Label')
219
    plt.tight_layout()
220
    plt.savefig('confusion_matrix.png', dpi=150)
221
    plt.close()
222

223
    return accuracy
224

225

226
if __name__ == '__main__':
227
    from models import MLP, ResNet18
228
    from torch.utils.data import TensorDataset, DataLoader
229

230
    # 构造模拟数据，快速验证训练流程而无需下载数据集
231
    batch_size = 16
232
    num_samples = 64
233
    x = torch.randn(num_samples, 3, 32, 32)
234
    y = torch.randint(0, 10, (num_samples,))
235

236
    dataset = TensorDataset(x, y)
237
    train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
238
    val_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
239
    test_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
240

241
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
242
    print(f"使用设备: {device}\n")
243

244
    print("=" * 50)
245
    print("测试 1: MLP Trainer 初始化与单 epoch 训练")
246
    print("=" * 50)
247
    mlp = MLP(num_classes=10)
248
    trainer_mlp = Trainer(mlp, train_loader, val_loader, device, model_name='MLP_test')
249
    train_loss, train_acc = trainer_mlp.train_epoch()
250
    print(f"训练 Loss: {train_loss:.4f}, Acc: {train_acc:.2f}%")
251
    val_loss, val_acc = trainer_mlp.validate()
252
    print(f"验证 Loss: {val_loss:.4f}, Acc: {val_acc:.2f}%")
253

254
    print("\n" + "=" * 50)
255
    print("测试 2: ResNet18 Trainer 初始化与单 epoch 训练")
256
    print("=" * 50)
257
    resnet = ResNet18(num_classes=10)
258
    trainer_resnet = Trainer(resnet, train_loader, val_loader, device, model_name='ResNet18_test')
259
    train_loss, train_acc = trainer_resnet.train_epoch()
260
    print(f"训练 Loss: {train_loss:.4f}, Acc: {train_acc:.2f}%")
261
    val_loss, val_acc = trainer_resnet.validate()
262
    print(f"验证 Loss: {val_loss:.4f}, Acc: {val_acc:.2f}%")
263

264
    print("\n" + "=" * 50)
265
    print("测试 3: MLP 完整训练流程（2 epochs）")
266
    print("=" * 50)
267
    mlp2 = MLP(num_classes=10)
268
    trainer_full = Trainer(mlp2, train_loader, val_loader, device, model_name='MLP_full')
269
    history = trainer_full.train(epochs=2, save_dir='./checkpoints_test')
270
    print(f"训练历史记录键: {list(history.keys())}")
271
    print(f"最佳验证准确率: {trainer_full.best_acc:.2f}%")
272

273
    print("\n" + "=" * 50)
274
    print("测试 4: test 函数（使用刚保存的模型）")
275
    print("=" * 50)
276
    model_path = './checkpoints_test/MLP_full_best.pth'
277
    mlp3 = MLP(num_classes=10)
278
    acc = test(mlp3, test_loader, device, model_path=model_path)
279
    print(f"test 函数返回准确率: {acc:.2f}%")
280

281
    import shutil
282

283
    if os.path.exists('./checkpoints_test'):
284
        shutil.rmtree('./checkpoints_test')
285
    print("\n所有测试通过！")

1
import torch
2
import os
3
from data_engineering import CIFAR10Data
4
from models import MLP, ResNet18
5
from train import Trainer, test
6

7

8
def main():
9
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
10
    print(f"使用设备: {device}")
11

12
    if not torch.cuda.is_available():
13
        print("警告: 未检测到CUDA，将使用CPU训练（速度较慢）")
14
        print("建议: 检查GPU驱动或安装CUDA版本的PyTorch")
15

16
    print("\n" + "=" * 50)
17
    print("步骤1: 数据工程")
18
    print("=" * 50)
19

20
    try:
21
        # Windows 平台 num_workers=0，避免多进程 DataLoader 报错
22
        data_manager = CIFAR10Data(batch_size=128, num_workers=0)
23
        train_loader, val_loader, test_loader = data_manager.get_loaders(val_ratio=0.1)
24
    except FileNotFoundError as e:
25
        print(e)
26
        print("\n请先准备CIFAR-10数据！")
27
        return
28

29
    print("\n" + "=" * 50)
30
    print("步骤2: MLP模型训练与测试")
31
    print("=" * 50)
32
    mlp = MLP(num_classes=10, dropout=0.3)
33
    mlp_params = sum(p.numel() for p in mlp.parameters()) / 1e6
34
    print(f"MLP参数量: {mlp_params:.2f}M")
35

36
    trainer_mlp = Trainer(
37
        mlp, train_loader, val_loader, device,
38
        lr=0.001, model_name='MLP'
39
    )
40
    trainer_mlp.train(epochs=10)
41

42
    # 训练结束后从 checkpoint 加载最优权重，在测试集评估
43
    test_acc_mlp = test(mlp, test_loader, device,
44
                       model_path='./checkpoints/MLP_best.pth')
45

46
    print("\n" + "=" * 50)
47
    print("步骤3: ResNet-18模型训练与测试")
48
    print("=" * 50)
49
    resnet = ResNet18(num_classes=10)
50
    resnet_params = sum(p.numel() for p in resnet.parameters()) / 1e6
51
    print(f"ResNet-18参数量: {resnet_params:.2f}M")
52

53
    trainer_resnet = Trainer(
54
        resnet, train_loader, val_loader, device,
55
        lr=0.001, model_name='ResNet18'
56
    )
57
    trainer_resnet.train(epochs=10)
58

59
    test_acc_resnet = test(resnet, test_loader, device,
60
                          model_path='./checkpoints/ResNet18_best.pth')
61

62
    print("\n" + "=" * 50)
63
    print("步骤4: 实验结果对比")
64
    print("=" * 50)
65
    print(f"{'模型':<15} {'参数量':<10} {'测试准确率':<10}")
66
    print("-" * 50)
67
    print(f"{'MLP':<15} {mlp_params:<10.2f} {test_acc_mlp:<10.2f}%")
68
    print(f"{'ResNet-18':<15} {resnet_params:<10.2f} {test_acc_resnet:<10.2f}%")
69
    print("=" * 50)
70

71
    with open('comparison_results.txt', 'w') as f:
72
        f.write("图像分类实验对比结果\n")
73
        f.write("=" * 50 + "\n")
74
        f.write(f"MLP: 参数量={mlp_params:.2f}M, 准确率={test_acc_mlp:.2f}%\n")
75
        f.write(f"ResNet-18: 参数量={resnet_params:.2f}M, 准确率={test_acc_resnet:.2f}%\n")
76

77
    print("\n结果已保存到 comparison_results.txt")
78

79

80
if __name__ == '__main__':
81
    main()

1
import torch
2
import matplotlib.pyplot as plt
3
from torchvision import transforms
4
from torch.utils.data import DataLoader, Subset
5
import numpy as np
6
import random
7
import os
8

9
from data_engineering import CIFAR10Data
10
from models import MLP, ResNet18
11

12

13
def denormalize(tensor, mean=(0.4914, 0.4822, 0.4465), std=(0.2023, 0.1994, 0.2010)):
14
    # 反归一化，把标准化后的张量还原到 [0,1] 区间用于显示
15
    for t, m, s in zip(tensor, mean, std):
16
        t.mul_(s).add_(m)
17
    return torch.clamp(tensor, 0, 1)
18

19

20
def visualize_comparison(num_samples=8, save_path='visualization_image.png'):
21
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
22
    print(f"使用设备: {device}")
23

24
    classes = ('plane', 'car', 'bird', 'cat', 'deer',
25
              'dog', 'frog', 'horse', 'ship', 'truck')
26

27
    # 只加载测试集用于抽样展示
28
    print("加载数据...")
29
    data_manager = CIFAR10Data(batch_size=1, num_workers=0)
30
    _, _, test_loader = data_manager.get_loaders(val_ratio=0.1)
31

32
    # 从测试集中随机抽取指定数量的样本
33
    test_size = len(test_loader.dataset)
34
    random_indices = random.sample(range(test_size), num_samples)
35

36
    subset = Subset(test_loader.dataset, random_indices)
37
    subset_loader = DataLoader(subset, batch_size=1, shuffle=False)
38

39
    print("加载模型...")
40

41
    mlp = MLP(num_classes=10).to(device)
42
    mlp.load_state_dict(torch.load('./checkpoints/MLP_best.pth', map_location=device))
43
    mlp.eval()
44

45
    resnet = ResNet18(num_classes=10).to(device)
46
    resnet.load_state_dict(torch.load('./checkpoints/ResNet18_best.pth', map_location=device))
47
    resnet.eval()
48

49
    images = []
50
    labels_true = []
51
    preds_mlp = []
52
    preds_resnet = []
53
    conf_mlp = []
54
    conf_resnet = []
55

56
    print("进行预测...")
57
    with torch.no_grad():
58
        for images_batch, labels_batch in subset_loader:
59
            images_batch = images_batch.to(device)
60

61
            # Softmax 将 logits 转为概率分布，取最大项作为预测类别与置信度
62
            outputs_mlp = mlp(images_batch)
63
            probs_mlp = torch.softmax(outputs_mlp, dim=1)
64
            pred_mlp = torch.argmax(probs_mlp, dim=1).item()
65
            conf_mlp_val = torch.max(probs_mlp).item()
66

67
            outputs_resnet = resnet(images_batch)
68
            probs_resnet = torch.softmax(outputs_resnet, dim=1)
69
            pred_resnet = torch.argmax(probs_resnet, dim=1).item()
70
            conf_resnet_val = torch.max(probs_resnet).item()
71

72
            img = denormalize(images_batch[0].cpu())
73
            images.append(img)
74
            labels_true.append(labels_batch.item())
75
            preds_mlp.append(pred_mlp)
76
            preds_resnet.append(pred_resnet)
77
            conf_mlp.append(conf_mlp_val)
78
            conf_resnet.append(conf_resnet_val)
79

80
    # 2x8 子图：上行展示 MLP 预测，下行展示 ResNet-18 预测
81
    fig, axes = plt.subplots(2, num_samples, figsize=(num_samples * 2, 4))
82

83
    for i in range(num_samples):
84
        ax_mlp = axes[0, i]
85
        ax_mlp.imshow(images[i].permute(1, 2, 0).numpy())
86
        ax_mlp.axis('off')
87

88
        pred_label = classes[preds_mlp[i]]
89
        is_correct = preds_mlp[i] == labels_true[i]
90

91
        # 绿色表示预测正确，红色表示预测错误
92
        color = 'green' if is_correct else 'red'
93
        text = f"{pred_label}\n{conf_mlp[i]:.2f}"
94

95
        ax_mlp.text(0.95, 0.95, text, transform=ax_mlp.transAxes,
96
                   fontsize=10, color=color, weight='bold',
97
                   ha='right', va='top',
98
                   bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
99

100
        ax_res = axes[1, i]
101
        ax_res.imshow(images[i].permute(1, 2, 0).numpy())
102
        ax_res.axis('off')
103

104
        pred_label = classes[preds_resnet[i]]
105
        is_correct = preds_resnet[i] == labels_true[i]
106

107
        color = 'green' if is_correct else 'red'
108
        text = f"{pred_label}\n{conf_resnet[i]:.2f}"
109

110
        ax_res.text(0.95, 0.95, text, transform=ax_res.transAxes,
111
                   fontsize=10, color=color, weight='bold',
112
                   ha='right', va='top',
113
                   bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
114

115
    fig.text(0.02, 0.75, 'MLP', ha='center', va='center',
116
            fontsize=14, weight='bold', rotation=90)
117
    fig.text(0.02, 0.25, 'ResNet-18', ha='center', va='center',
118
            fontsize=14, weight='bold', rotation=90)
119

120
    plt.suptitle('Image Classification: MLP vs ResNet-18\n(Green=Correct, Red=Wrong)',
121
                fontsize=14, weight='bold')
122
    plt.tight_layout(rect=[0.03, 0, 1, 0.95])
123
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
124
    plt.close()
125

126
    print(f"\n✓ 可视化结果已保存: {save_path}")
127

128
    # 统计两个模型在这批样本上的正确率
129
    correct_mlp = sum(1 for p, t in zip(preds_mlp, labels_true) if p == t)
130
    correct_resnet = sum(1 for p, t in zip(preds_resnet, labels_true) if p == t)
131

132
    print(f"\n统计（{num_samples}个样本）:")
133
    print(f"  MLP正确率: {correct_mlp}/{num_samples} ({100*correct_mlp/num_samples:.1f}%)")
134
    print(f"  ResNet正确率: {correct_resnet}/{num_samples} ({100*correct_resnet/num_samples:.1f}%)")
135

136
    return save_path
137

138

139
if __name__ == '__main__':
140
    # 权重文件缺失时提示先运行 main.py
141
    if not os.path.exists('./checkpoints/MLP_best.pth'):
142
        print("错误：找不到 MLP_best.pth，请先运行 main.py 训练模型")
143
        exit(1)
144

145
    if not os.path.exists('./checkpoints/ResNet18_best.pth'):
146
        print("错误：找不到 ResNet18_best.pth，请先运行 main.py 训练模型")
147
        exit(1)
148

149
    visualize_comparison(num_samples=8)