当前位置：首页 > news >正文

深度学习---pytorch搭建深度学习模型（附带图片五分类实例）

news 2025/4/30 6:12:58

一、PyTorch搭建深度学习模型流程

1. 环境准备

安装PyTorch及相关库：

pip install torch torchvision numpy matplotlib

2. 数据准备

数据集加载：使用内置数据集（如CIFAR-10）或自定义数据集。
数据预处理：包括归一化、数据增强（随机翻转、旋转等）。
数据划分：将数据集分为训练集、验证集和测试集。
数据加载器：使用DataLoader实现批量加载。

3. 模型构建

网络结构：通过继承nn.Module定义模型，使用卷积层、池化层、全连接层等。
激活函数：如ReLU、Sigmoid。
正则化：Dropout层、BatchNorm层。
损失函数：如交叉熵损失（CrossEntropyLoss）。
优化器：如Adam、SGD。

4. 训练流程

前向传播：计算模型输出。
损失计算：根据预测和标签计算损失。
反向传播：通过loss.backward()计算梯度。
参数更新：优化器通过optimizer.step()更新权重。
训练循环：多轮迭代训练，记录训练损失和准确率。

5. 验证与测试

模型评估模式：使用model.eval()关闭Dropout和BatchNorm。
禁用梯度计算：使用torch.no_grad()加速推理。
指标计算：如准确率、F1分数。

6. 模型保存与加载

保存模型参数：torch.save(model.state_dict(), "model.pth")
加载模型参数：model.load_state_dict(torch.load("model.pth"))

7. 高级功能

学习率调度：动态调整学习率（如StepLR）。
自定义层/损失函数：继承nn.Module或autograd.Function。
多GPU训练：使用DataParallel或DistributedDataParallel。

二、图片五分类实例（基于CIFAR-10子集）

1. 数据准备

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Subset
import numpy as np# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 数据预处理与增强
transform_train = transforms.Compose([transforms.RandomHorizontalFlip(),transforms.RandomRotation(10),transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])transform_test = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])# 加载CIFAR-10并过滤前5类
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True)
test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True)# 过滤出前5类（类别0-4）
train_idx = np.isin(train_dataset.targets, [0,1,2,3,4])
test_idx = np.isin(test_dataset.targets, [0,1,2,3,4])train_dataset = Subset(train_dataset, np.where(train_idx)[0])
test_dataset = Subset(test_dataset, np.where(test_idx)[0])# 应用Transform
train_dataset.dataset.transform = transform_train
test_dataset.dataset.transform = transform_test# 划分训练集和验证集
train_size = int(0.8 * len(train_dataset))
val_size = len(train_dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])# 创建数据加载器
batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2)

2. 模型定义

class CNN(nn.Module):def __init__(self, num_classes=5):super(CNN, self).__init__()self.features = nn.Sequential(nn.Conv2d(3, 16, kernel_size=3, padding=1),nn.ReLU(),nn.MaxPool2d(2, 2),nn.Conv2d(16, 32, kernel_size=3, padding=1),nn.ReLU(),nn.MaxPool2d(2, 2),)self.classifier = nn.Sequential(nn.Linear(32 * 8 * 8, 256),nn.ReLU(),nn.Dropout(0.5),nn.Linear(256, num_classes),)def forward(self, x):x = self.features(x)x = x.view(x.size(0), -1)x = self.classifier(x)return xmodel = CNN(num_classes=5).to(device)

3. 训练配置

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

4. 训练循环

num_epochs = 10
best_val_acc = 0.0for epoch in range(num_epochs):# 训练阶段model.train()train_loss, train_correct, total = 0, 0, 0for 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()train_loss += loss.item() * images.size(0)_, predicted = outputs.max(1)total += labels.size(0)train_correct += predicted.eq(labels).sum().item()# 验证阶段model.eval()val_loss, val_correct, val_total = 0, 0, 0with 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 = outputs.max(1)val_total += labels.size(0)val_correct += predicted.eq(labels).sum().item()# 打印结果print(f"Epoch {epoch+1}/{num_epochs}")print(f"Train Loss: {train_loss/total:.4f}, Acc: {train_correct/total:.4f}")print(f"Val Loss: {val_loss/val_total:.4f}, Acc: {val_correct/val_total:.4f}\n")# 保存最佳模型if val_correct/val_total > best_val_acc:best_val_acc = val_correct/val_totaltorch.save(model.state_dict(), "best_model.pth")

5. 测试模型

model.load_state_dict(torch.load("best_model.pth"))
model.eval()
test_correct, test_total = 0, 0
with torch.no_grad():for images, labels in test_loader:images, labels = images.to(device), labels.to(device)outputs = model(images)_, predicted = outputs.max(1)test_total += labels.size(0)test_correct += predicted.eq(labels).sum().item()print(f"Test Accuracy: {test_correct/test_total:.4f}")