Upload New File

5-深度学习/5.6-作业/train.py

+#  此训练脚本更改了第7行 60行
+#  此训练脚本更改了第7行 60行
+from torchvision import transforms
+from torch.utils.data import DataLoader
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from tqdm import tqdm
+from model import SimpleCNN,BetterCNN
+from torchvision.datasets import ImageFolder
+from my_dataset import CatDogDataset,CatDogDataset_new
+
+# 定义训练集的图像预处理（transform）操作
+# transforms.RandomRotation(90) 随机旋转图像，用于数据增强
+# transforms.RandomHorizontalFlip(p=0.5)  以50%概率水平翻转图像，用于数据增强
+train_transform = transforms.Compose([
+    transforms.Resize((256, 256)),                  # 将图片缩放到统一大小 256x256，方便批量训练
+    transforms.RandomCrop(224),                     #  随机裁剪一个224x224的区域，增强模型的鲁棒性
+    transforms.ToTensor(),                          # 将 PIL 图像或 ndarray 转换为 Tensor，并归一化到 [0,1]
+    transforms.Normalize([0.5, 0.5, 0.5],           # 对三个通道分别做标准化：x=(x-mean)/std,这样处理后像素值范围变为 [-1, 1]
+                         [0.5, 0.5, 0.5])           # 让数据分布更均匀，减轻训练时梯度更新的震荡，有助于模型更快收敛
+])
+
+val_transform = transforms.Compose([
+    transforms.Resize((256, 256)),           # 统一调整图像大小，保证输入尺寸一致
+    transforms.CenterCrop(224),              # 中心裁剪出224x224大小区域，保证验证时稳定性
+    transforms.ToTensor(),                   # 转为Tensor，同时将像素值归一化到0~1之间
+    transforms.Normalize([0.5, 0.5, 0.5],    # 标准化，三个通道均值
+                         [0.5, 0.5, 0.5])    # 标准差，三个通道均方差
+])
+
+# ImageFolder为torchvision.datasets提供的一个图像分类数据集加载器，专门用于处理按文件夹组织的分类数据
+
+# 创建训练集对象，加载路径data2\train下的图像，应用训练时的数据增强和预处理（如随机裁剪、翻转等）
+train_dataset = CatDogDataset_new(r"data\train", transform=train_transform) 
+# 创建验证集对象，加载路径data2\val下的图像，应用验证时的固定预处理（如中心裁剪），保证验证结果稳定
+val_dataset = CatDogDataset_new(r"data\val", transform=val_transform)     
+
+# 创建数据加载器
+# DataLoader就是用来把你的数据集“打包好，一批一批地送进模型里”的工具
+train_loader = DataLoader(
+    train_dataset,       # 传入训练数据集对象
+    batch_size=32,       # 每个批次加载32张图像
+    shuffle=True,        # 打乱数据顺序，有利于模型更好地泛化
+    num_workers=0        # 用于数据加载的子进程数量，0表示在主进程加载数据
+)
+
+val_loader = DataLoader(
+    val_dataset,         # 传入验证数据集对象
+    batch_size=32,       # 验证时也用相同批大小
+    shuffle=False,       # 验证时不打乱数据，保证结果稳定可复现
+    num_workers=0        # 同样在主进程加载数据
+)
+
+
+# 设备选择
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# 实例化模型、损失函数、优化器
+# model = SimpleCNN().to(device)
+model = BetterCNN().to(device)  # 猫狗分类是两个类别
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+# 训练函数
+def train_model(model, train_loader, val_loader, criterion, optimizer, epochs=10):
+    # 训练多个轮次（epoch）
+    for epoch in range(epochs):
+        model.train()  # 设置模型为训练模式，启用dropout等
+        train_loss, correct, total = 0, 0, 0  # 初始化训练损失、正确预测数量、总样本数
+        # 遍历训练集的所有批次
+        for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}"):
+            # 将数据移动到GPU（如果使用GPU训练）
+            images, labels = images.to(device), labels.to(device)
+            # 前向传播：将图像输入模型，得到预测结果
+            outputs = model(images)
+            # 计算损失（预测结果 vs 实际标签）
+            loss = criterion(outputs, labels)
+            # 清空上一轮的梯度信息
+            optimizer.zero_grad()
+            # 反向传播：计算梯度
+            loss.backward()
+            # 更新模型参数
+            optimizer.step()
+            # 累加损失：loss.item() 是当前batch的平均损失
+            train_loss += loss.item() * images.size(0)  # 乘上batch大小是为了计算总损失
+            # 获取预测结果：取输出中最大值所在的类别索引
+            _, predicted = torch.max(outputs, 1)
+            # 累加总样本数和预测正确的样本数
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+        # 计算本轮训练集上的准确率
+        acc = 100 * correct / total
+        print(f"Train Loss: {train_loss/total:.4f}, Accuracy: {acc:.2f}%")
+
+        # ------------------ 验证阶段 ------------------
+        model.eval()  # 切换到评估模式（关闭dropout等）
+        val_loss, correct, total = 0, 0, 0
+
+        # 在评估阶段不需要计算梯度，加速并节省显存
+        with torch.no_grad():
+            for images, labels in val_loader:
+                images, labels = images.to(device), labels.to(device)
+                # 前向传播得到预测
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+                # 累加验证损失
+                val_loss += loss.item() * images.size(0)
+                # 获取预测类别
+                _, predicted = torch.max(outputs, 1)
+                # 累加验证集中总数与正确数
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+        # 计算验证准确率
+        val_acc = 100 * correct / total
+        print(f"Val Loss: {val_loss/total:.4f}, Accuracy: {val_acc:.2f}%\n")
+# 调用训练函数：传入模型、训练/验证数据加载器、损失函数、优化器和训练轮次
+train_model(model, train_loader, val_loader, criterion, optimizer, epochs=10)
+# 将训练好的模型参数保存到本地文件，方便之后加载使用或部署
+torch.save(model.state_dict(), "cat_dog_cnn_new.pth")