1
import torch
2
import torch.nn as nn
3
import torch.optim as optim
4
from torchvision import datasets, transforms
5
from torch.utils.data import DataLoader
6
import matplotlib.pyplot as plt
7
import numpy as np
8
import os
9
import shutil
10

11
# ==========================================
12
# 1. 配置与环境清理 (Setup & Cleaning)
13
# ==========================================
14

15
# 检查 GPU 并定义 device
16
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
print(f"当前使用的计算设备: {device}")
18

19
# --- V3  超参数设置 ---
20
BATCH_SIZE = 128
21
EPOCHS = 100
22
HEIGHT = 150
23
WIDTH = 150
24
LEARNING_RATE = 0.001
25
WEIGHT_DECAY = 5e-5 # V3 优化：降低 L2 权重衰减系数
26

27
# 数据集路径 (依赖本地已有的文件)
28
DATASET_DIR = 'cats_and_dogs_filtered'
29
train_dir = os.path.join(DATASET_DIR, 'train')
30
validation_dir = os.path.join(DATASET_DIR, 'validation')
31

32

33
# --- 环境清理 (清理 .ipynb_checkpoints) ---
34

35
def clean_checkpoints(directory):
36
    """清理 .ipynb_checkpoints 防止 FileNotFoundError"""
37
    if not os.path.exists(directory):
38
        return
39
    for root, dirs, files in os.walk(directory):
40
        for d in dirs:
41
            if d == ".ipynb_checkpoints":
42
                path = os.path.join(root, d)
43
                print(f"删除干扰目录: {path}")
44
                shutil.rmtree(path)
45

46
# 在数据加载前执行清理
47
clean_checkpoints(DATASET_DIR)
48

49

50
# ==========================================
51
# 2. 数据加载与预处理 (Data Loading) - 增加数据增强
52
# ==========================================
53

54
# 定义转换
55
data_transforms = {
56
    'train': transforms.Compose([
57
        transforms.Resize((HEIGHT, WIDTH)),
58
        transforms.RandomHorizontalFlip(),
59
        transforms.RandomRotation(15),
60
        transforms.ToTensor(),
61
    ]),
62
    'val': transforms.Compose([
63
        transforms.Resize((HEIGHT, WIDTH)),
64
        transforms.ToTensor(),
65
    ]),
66
}
67

68
# 加载数据集
69
try:
70
    train_dataset = datasets.ImageFolder(train_dir, transform=data_transforms['train'])
71
    val_dataset = datasets.ImageFolder(validation_dir, transform=data_transforms['val'])
72

73
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
74
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
75

76
    print(f"数据加载成功: 训练集 {len(train_dataset)} 张, 验证集 {len(val_dataset)} 张")
77
except FileNotFoundError as e:
78
    # 如果到这里还出错，则提示用户检查路径
79
    print(f"\nFATAL ERROR: 数据加载失败，请检查 '{train_dir}' 路径是否正确。")
80
    print("错误详情: 请确保 'cats_and_dogs_filtered/train' 目录已存在并包含数据。")
81
    # 强制终止，避免 NameError
82
    raise SystemExit("程序终止：无法找到训练数据。")

1
# ==========================================
2
# 3. 数据可视化预览 (Data Visualization)
3
# ==========================================
4

5
def visualize_data(loader):
6
    """展示一批数据以确认预处理效果"""
7
    dataiter = iter(loader)
8
    images, labels = next(dataiter)
9

10
    # 调整图像以适应 matplotlib 显示 (C, H, W) -> (H, W, C)
11
    npimg = images[0].numpy().transpose((1, 2, 0))
12

13
    plt.figure(figsize=(4, 4))
14
    # PyTorch Tensors are [0, 1], so no need to normalize for imshow
15
    plt.imshow(npimg)
16
    title = train_dataset.classes[labels[0].item()]
17
    plt.title(title)
18
    plt.axis('off')
19
    plt.show()

1
# ==========================================
2
# 4. 模型构建 (Model Architecture) - V3 增强 Dropout (p=0.6)
3
# ==========================================
4

5
class CatDogCNN(nn.Module):
6
    def __init__(self):
7
        super(CatDogCNN, self).__init__()
8

9
        # 卷积层 (BN 结构: Conv -> BN -> ReLU -> Pool)
10
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
11
        self.bn1 = nn.BatchNorm2d(16)
12
        self.pool = nn.MaxPool2d(2, 2)
13

14
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
15
        self.bn2 = nn.BatchNorm2d(32)
16

17
        self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
18
        self.bn3 = nn.BatchNorm2d(64)
19

20
        self.relu = nn.ReLU()
21
        # V3 优化：增强 Dropout 率
22
        self.dropout = nn.Dropout(p=0.6)
23

24
        # 64 * 18 * 18 是经过三次 MaxPool2d(2, 2) 后特征图的展平尺寸 (150/2/2/2 = 18.75 -> 18)
25
        self.fc1 = nn.Linear(64 * 18 * 18, 512)
26
        self.fc2 = nn.Linear(512, 1)
27

28
    def forward(self, x):
29
        # 模块 1
30
        x = self.conv1(x); x = self.bn1(x); x = self.relu(x); x = self.pool(x)
31

32
        # 模块 2
33
        x = self.conv2(x); x = self.bn2(x); x = self.relu(x); x = self.pool(x)
34

35
        # 模块 3
36
        x = self.conv3(x); x = self.bn3(x); x = self.relu(x); x = self.pool(x)
37

38
        # 展平
39
        x = x.view(x.size(0), -1)
40

41
        # 全连接层 with Dropout
42
        x = self.dropout(x)
43
        x = self.relu(self.fc1(x))
44
        x = self.dropout(x)
45

46
        x = self.fc2(x)
47
        return x
48

49
model = CatDogCNN().to(device)

1
# ==========================================
2
# 5. 训练配置 (Training Setup) - V3 优化 Scheduler Patience
3
# ==========================================
4

5
criterion = nn.BCEWithLogitsLoss()
6
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
7

8
# V3 优化：学习率调度器，增加 patience
9
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
10
    optimizer,
11
    mode='min',
12
    factor=0.2,
13
    patience=8,          # V3: 从 5 增加到 8
14
    min_lr=1e-6,
15
    verbose=True
16
)
17

18
# Early Stopping 配置 (保持不变)
19
patience = 15
20
min_val_loss = np.Inf
21
patience_counter = 0
22
best_epoch = 0
23
history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}
24

25
# ==========================================
26
# 6. 训练循环 (Training Loop)
27
# ==========================================
28

29
print(f"开始训练，共 {EPOCHS} 个 Epochs...")
30

31
for epoch in range(EPOCHS):
32
    # --- 训练阶段 ---
33
    model.train()
34
    running_loss = 0.0
35
    correct_train = 0
36
    total_train = 0
37

38
    for inputs, labels in train_loader:
39
        inputs, labels = inputs.to(device), labels.float().unsqueeze(1).to(device)
40

41
        optimizer.zero_grad()
42
        outputs = model(inputs)
43
        loss = criterion(outputs, labels)
44
        loss.backward()
45
        optimizer.step()
46

47
        running_loss += loss.item() * inputs.size(0)
48
        # sigmoid -> threshold for accuracy calculation
49
        predicted = (torch.sigmoid(outputs) > 0.5).float()
50
        correct_train += (predicted == labels).sum().item()
51
        total_train += labels.size(0)
52

53
    epoch_loss = running_loss / len(train_dataset)
54
    epoch_acc = correct_train / total_train
55

56
    # --- 验证阶段 ---
57
    model.eval()
58
    val_running_loss = 0.0
59
    correct_val = 0
60
    total_val = 0
61

62
    with torch.no_grad():
63
        for inputs, labels in val_loader:
64
            inputs, labels = inputs.to(device), labels.float().unsqueeze(1).to(device)
65
            outputs = model(inputs)
66
            loss = criterion(outputs, labels)
67

68
            val_running_loss += loss.item() * inputs.size(0)
69
            predicted = (torch.sigmoid(outputs) > 0.5).float()
70
            correct_val += (predicted == labels).sum().item()
71
            total_val += labels.size(0)
72

73
    val_loss = val_running_loss / len(val_dataset)
74
    val_acc = correct_val / total_val
75

76
    # 调用调度器并打印 LR
77
    scheduler.step(val_loss)
78
    current_lr = optimizer.param_groups[0]['lr']
79

80
    # --- 记录历史 ---
81
    history['train_loss'].append(epoch_loss)
82
    history['train_acc'].append(epoch_acc)
83
    history['val_loss'].append(val_loss)
84
    history['val_acc'].append(val_acc)
85

86
    # --- 日志打印 ---
87
    print(f"Epoch [{epoch+1}/{EPOCHS}] "
88
          f"Train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f} | "
89
          f"Val Loss: {val_loss:.4f} Acc: {val_acc:.4f} | LR: {current_lr:.6f}")
90

91
    # --- Early Stopping 逻辑 ---
92
    if val_loss < min_val_loss:
93
        min_val_loss = val_loss
94
        best_epoch = epoch + 1
95
        patience_counter = 0
96
        print(f"   验证损失下降，当前最佳模型保存在 Epoch {best_epoch}")
97
    else:
98
        patience_counter += 1
99
        print(f"   验证损失未下降，耐心计数: {patience_counter}/{patience}")
100
        if patience_counter >= patience:
101
            print(f"Early Stopping 触发！最佳 Epoch 为 {best_epoch}，最小验证损失 {min_val_loss:.4f}。")
102
            break
103

104
print("训练完成！")

1
# ==========================================
2
# 7. 结果可视化 (Result Visualization)
3
# ==========================================
4

5
epochs_range = range(len(history['train_acc']))
6

7
plt.figure(figsize=(12, 5))
8

9
# 准确率曲线
10
plt.subplot(1, 2, 1)
11
plt.plot(epochs_range, history['train_acc'], label='Training Accuracy')
12
plt.plot(epochs_range, history['val_acc'], label='Validation Accuracy')
13
plt.axvline(best_epoch - 1, color='r', linestyle='--', label=f'Best Epoch ({best_epoch})')
14
plt.legend(loc='lower right')
15
plt.title('Training and Validation Accuracy')
16

17
# 损失曲线
18
plt.subplot(1, 2, 2)
19
plt.plot(epochs_range, history['train_loss'], label='Training Loss')
20
plt.plot(epochs_range, history['val_loss'], label='Validation Loss')
21
plt.axvline(best_epoch - 1, color='r', linestyle='--', label=f'Best Epoch ({best_epoch})')
22
plt.legend(loc='upper right')
23
plt.title('Training and Validation Loss')
24

25
plt.show()