深度学习篇---分类任务图像预处理&模型训练

Ronin-Lotus2025-03-20 22:22

文章目录

前言

本文简单介绍了pytoch、paddlepaddle框架下的分类任务的图像预处理、模型训练以及模型保存的流程。

一、Pytorch

python 复制代码 
import os
import cv2
import numpy as np
import torch
from torch import nn, optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.models import resnet18

#-------------------- 数据读取与预处理 --------------------
class CustomDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = data_dir
        self.classes = os.listdir(data_dir)
        self.image_paths = []
        self.labels = []

        # 遍历目录，获取所有图像路径和标签
        for label, cls in enumerate(self.classes):
            cls_dir = os.path.join(data_dir, cls)
            for img_name in os.listdir(cls_dir):
                self.image_paths.append(os.path.join(cls_dir, img_name))
                self.labels.append(label)

        self.transform = transform

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = cv2.imread(img_path)  # 使用OpenCV读取图像
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转为RGB格式

        if self.transform:
            image = self.transform(image)
        
        label = self.labels[idx]
        return image, label

#定义数据增强和归一化
train_transform = transforms.Compose([
    transforms.ToPILImage(),  # OpenCV图像转PIL格式
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

val_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

#创建Dataset和DataLoader
train_dataset = CustomDataset('data/train', transform=train_transform)
val_dataset = CustomDataset('data/val', transform=val_transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

#-------------------- 模型定义 --------------------
class CustomModel(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.backbone = resnet18(pretrained=True)
        self.backbone.fc = nn.Linear(self.backbone.fc.in_features, num_classes)

    def forward(self, x):
        return self.backbone(x)

model = CustomModel(num_classes=2)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

#-------------------- 训练配置 --------------------
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)

#-------------------- 训练循环 --------------------
for epoch in range(10):
    model.train()
    train_loss = 0.0

    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    # 验证
    model.eval()
    val_loss = 0.0
    correct = 0
    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            val_loss += criterion(outputs, labels).item()
            _, preds = torch.max(outputs, 1)
            correct += (preds == labels).sum().item()

    print(f'Epoch {epoch+1}, Train Loss: {train_loss/len(train_loader):.4f}, Val Acc: {correct/len(val_dataset):.4f}')

#-------------------- 模型导出 --------------------
#保存PyTorch模型权重
torch.save(model.state_dict(), 'model.pth')

#导出为ONNX格式（可选）
dummy_input = torch.randn(1, 3, 224, 224).to(device)
torch.onnx.export(model, dummy_input, 'model.onnx', input_names=['input'], output_names=['output'])

1. 自定义数据集类 CustomDataset

python 复制代码 
class CustomDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        # 初始化数据集路径和标签
        self.data_dir = data_dir
        self.classes = os.listdir(data_dir)  # 获取类别文件夹（如class1, class2）
        self.image_paths = []  # 存储所有图像路径
        self.labels = []        # 存储对应标签
        # 遍历子文件夹，构建路径和标签的映射
        for label, cls in enumerate(self.classes):
            cls_dir = os.path.join(data_dir, cls)
            for img_name in os.listdir(cls_dir):
                self.image_paths.append(os.path.join(cls_dir, img_name))
                self.labels.append(label)
        self.transform = transform  # 数据增强/归一化操作

    def __len__(self):
        return len(self.image_paths)  # 返回数据集总样本数

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = cv2.imread(img_path)  # 用OpenCV读取图像（BGR格式）
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转为RGB格式（适配模型输入）
        if self.transform:
            image = self.transform(image)  # 应用预处理
        label = self.labels[idx]
        return image, label  # 返回张量和标签

关键点

init

init : 递归遍历文件夹，生成图像路径与标签的映射。

getitem

getitem : 动态加载图像 ，OpenCV读取后需转为RGB（因PyTorch模型默认处理RGB）。

transform

transform: 接收外部定义的数据增强操作（如归一化、翻转）。

2. 数据预处理 transforms.Compose

python 复制代码 
train_transform = transforms.Compose([
    transforms.ToPILImage(),  # 将numpy数组或OpenCV图像转为PIL格式（后续操作需要）
    transforms.Resize((224, 224)),  # 调整图像尺寸为224x224
    transforms.RandomHorizontalFlip(),  # 随机水平翻转（概率默认0.5）
    transforms.RandomRotation(15),      # 随机旋转（-15度到+15度）
    transforms.ToTensor(),  # 转为Tensor格式（维度变为CxHxW，像素值范围[0,1]）
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # ImageNet统计量归一化
])

关键点

ToPILmage

ToPILImage: PyTorch的 transforms 多数操作基于PIL图像，需先转换格式。

Normalize

Normalize: 使用ImageNet的均值和标准差，适配预训练模型输入分布。

验证集无需数据增强

验证集无需数据增强：仅保留尺寸调整和归一化。

3. 数据加载器 DataLoader

python 复制代码 
train_loader = DataLoader(
    train_dataset, 
    batch_size=32,   # 每个批次的样本数
    shuffle=True,    # 打乱训练数据顺序（防止过拟合）
    num_workers=4    # 使用4个子进程加载数据（加速数据读取）
)

作用

作用：将数据集按批次加载，支持多进程加速。

参数

num_workers: 根据CPU核心数调整，避免过高导致内存溢出。

4. 模型定义 CustomModel

python 复制代码 
class CustomModel(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.backbone = resnet18(pretrained=True)  # 加载预训练ResNet18
        # 修改全连接层，适配自定义类别数
        self.backbone.fc = nn.Linear(self.backbone.fc.in_features, num_classes)

    def forward(self, x):
        return self.backbone(x)

关键点

预训练权重

预训练权重：pretrained=True 加载ImageNet预训练参数，加速收敛。

修改全连接层

修改全连接层：原始ResNet输出1000类 ，需替换为任务实际类别数（如2类）。

5. 训练循环

python 复制代码 
model.to(device)  # 将模型移至GPU（若可用）
criterion = nn.CrossEntropyLoss()  # 分类任务常用损失函数
optimizer = optim.Adam(model.parameters(), lr=1e-4)  # Adam优化器

for epoch in range(10):
    model.train()  # 设置为训练模式（启用BatchNorm和Dropout）
    train_loss = 0.0
    for images, labels in train_loader:
        images = images.to(device)  # 数据移至GPU
        labels = labels.to(device)
        optimizer.zero_grad()  # 清空梯度（防止累积）
        outputs = model(images)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()          # 反向传播计算梯度
        optimizer.step()         # 更新权重
        train_loss += loss.item()  # 累加损失

    # 验证阶段
    model.eval()  # 设置为评估模式（关闭BatchNorm和Dropout）
    with torch.no_grad():  # 禁用梯度计算（节省内存）
        for images, labels in val_loader:
            # ...（类似训练循环，但只计算损失和精度）

关键点

model.train()和model.evel()

model.train() 和 model.eval(): 控制模型训练/评估模式，影响某些层的行为。

optimizer.zero_grad()

optimizer.zero_grad(): 必须清空梯度 ，否则梯度会累积导致训练异常。

6. 模型导出

python 复制代码 
#保存PyTorch权重
torch.save(model.state_dict(), 'model.pth')  # 仅保存模型参数（轻量）

#导出为ONNX格式（跨框架部署）
dummy_input = torch.randn(1, 3, 224, 224).to(device)  # 生成虚拟输入
torch.onnx.export(
    model, 
    dummy_input, 
    'model.onnx', 
    input_names=['input'],  # 输入/输出名称（部署时标识）
    output_names=['output']
)

关键点

ONNX导出

ONNX导出：需定义输入张量的形状（如 1,3,224,224） ，便于后续部署到不同框架。

二、Paddlepaddle

python 复制代码 
import os
import cv2
import numpy as np
import paddle
from paddle.io import Dataset, DataLoader
from paddle.vision import transforms
from paddle.vision.models import resnet18

#-------------------- 数据读取与预处理 --------------------
class CustomDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        super().__init__()
        self.data_dir = data_dir
        self.classes = os.listdir(data_dir)
        self.image_paths = []
        self.labels = []

        for label, cls in enumerate(self.classes):
            cls_dir = os.path.join(data_dir, cls)
            for img_name in os.listdir(cls_dir):
                self.image_paths.append(os.path.join(cls_dir, img_name))
                self.labels.append(label)

        self.transform = transform

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = cv2.imread(img_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        if self.transform:
            image = self.transform(image)
        
        label = self.labels[idx]
        return image, label

    def __len__(self):
        return len(self.image_paths)

#数据增强与归一化
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

#创建DataLoader
train_dataset = CustomDataset('data/train', transform=train_transform)
val_dataset = CustomDataset('data/val', transform=val_transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

#-------------------- 模型定义 --------------------
class CustomModel(paddle.nn.Layer):
    def __init__(self, num_classes):
        super().__init__()
        self.backbone = resnet18(pretrained=True)
        self.backbone.fc = paddle.nn.Linear(self.backbone.fc.weight.shape[0], num_classes)

    def forward(self, x):
        return self.backbone(x)

model = CustomModel(num_classes=2)
model = paddle.Model(model)

#-------------------- 训练配置 --------------------
optimizer = paddle.optimizer.Adam(parameters=model.parameters(), learning_rate=1e-4)
loss_fn = paddle.nn.CrossEntropyLoss()
metric = paddle.metric.Accuracy()

#-------------------- 训练循环 --------------------
model.prepare(optimizer, loss_fn, metric)
model.fit(train_loader, val_loader, epochs=10, verbose=1)

#-------------------- 模型导出 --------------------
#保存Paddle模型权重
paddle.save(model.state_dict(), 'model.pdparams')

#导出为推理模型（静态图）
input_spec = paddle.static.InputSpec(shape=[None, 3, 224, 224], dtype='float32', name='input')
paddle.jit.save(model.network, 'inference_model', input_spec=[input_spec])

1. 数据集类 CustomDataset

与PyTorch实现逻辑一致，但继承自 paddle.io.Dataset，代码结构相同。

2. 数据预处理 transforms

python 复制代码 
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 直接支持numpy数组，无需转为PIL
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ToTensor(),  # 转为Tensor（维度为CxHxW）
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

差异点

差异点：Paddle的 transforms 直接支持numpy输入，无需 ToPILImage 转换。

3. 模型定义 CustomModel

python 复制代码 
class CustomModel(paddle.nn.Layer):
    def __init__(self, num_classes):
        super().__init__()
        self.backbone = resnet18(pretrained=True)
        # 修改全连接层（需手动获取输入维度）
        self.backbone.fc = paddle.nn.Linear(self.backbone.fc.weight.shape[0], num_classes)
    
    def forward(self, x):
        return self.backbone(x)

差异点

差异点：Paddle的 resnet18 全连接层属性名称为 fc ，与PyTorch一致，但需通过 weight.shape[0] 获取输入维度。

4. 训练配置与执行

python 复制代码 
model = paddle.Model(model)  # 高层API封装
model.prepare(
    optimizer=paddle.optimizer.Adam(parameters=model.parameters(), learning_rate=1e-4),
    loss=paddle.nn.CrossEntropyLoss(),
    metrics=paddle.metric.Accuracy()  # 内置评估指标
)
model.fit(train_loader, val_loader, epochs=10, verbose=1)  # 自动执行训练和验证

关键点

高层API

高层API: PaddlePaddle的 Model 类封装了训练循环，简化代码。

model.prepare(）

model.prepare(): 绑定优化器、损失函数和评估指标。

model.fit()

model.fit(): 自动执行多轮训练和验证 ，无需手动编写循环。

5. 模型导出

python 复制代码 
#保存权重
paddle.save(model.state_dict(), 'model.pdparams')  # 类似PyTorch的.pth

#导出静态图模型（用于部署）
input_spec = paddle.static.InputSpec(
    shape=[None, 3, 224, 224],  # 支持动态batch（None）
    dtype='float32', 
    name='input'
)
paddle.jit.save(model.network, 'inference_model', input_spec=[input_spec])

关键点

静态图导出

静态图导出: Paddle的 paddle.jit.save 生成部署专用模型，支持推理优化。

三、核心差异总结

功能 PyTorch PaddlePaddle

数据读取 需手动转为PIL格式 直接支持numpy数组

模型定义 nn.Module + 手动训练循环 paddle.nn.Layer + 高层API Model

训练循环 手动编写循环（灵活） 封装为 model.fit()（简洁）

模型导出 支持 .pth 和 ONNX 支持 .pdparams 和静态图模型

四、常见问题

  1. 为何使用OpenCV而非PIL读取图像？

OpenCV读取速度快 ，且支持更多格式（如16位图像） ，但需注意BGR转RGB。

  1. num_workers 设置多少合适？

通常设为CPU核心数的2~4倍，但需根据内存调整。设为0时禁用多进程。

  1. 验证集是否需要 shuffle=False？

是的，验证集需保持顺序一致，确保评估结果稳定。

  1. 如何适配其他模型（如Vision Transformer）？

修改 CustomModel 中的 backbone，替换为对应模型结构。

相关推荐
老歌老听老掉牙几秒前
SymPy 矩阵到 NumPy 数组的全面转换指南
python·线性代数·矩阵·numpy·sympy
CoovallyAIHub1 分钟前
方案 | 光伏清洁机器人系统详细技术实施方案
深度学习·算法·计算机视觉
星期天要睡觉3 分钟前
机器学习——CountVectorizer将文本集合转换为 基于词频的特征矩阵
人工智能·机器学习·矩阵
lxmyzzs5 分钟前
【图像算法 - 14】精准识别路面墙体裂缝：基于YOLO12与OpenCV的实例分割智能检测实战（附完整代码）
人工智能·opencv·算法·计算机视觉·裂缝检测·yolo12
站大爷IP6 分钟前
Python条件判断：从基础到进阶的实用指南
python
赛博郎中10 分钟前
pygame小游戏飞机大战_8继承精灵玩家优化
python·pygame
什么都想学的阿超14 分钟前
【大语言模型 01】注意力机制数学推导：从零实现Self-Attention
人工智能·语言模型·自然语言处理
William一直在路上22 分钟前
Python数据类型转换详解：从基础到实践
开发语言·python
trayvontang1 小时前
Python虚拟环境与包管理工具(uv、Conda)
python·conda·uv·虚拟环境·miniconda·miniforge
热门推荐
01UV安装并设置国内源02KGG转MP3工具|非KGM文件|解密音频03【2025.08.06最新版】Android Studio下载、安装及配置记录（自动下载sdk）04Qwen3-Coder 快速上手教程 | Qwen Code + Claude Code052025最新国内服务器可用docker源仓库地址大全（2025年8月更新）06蜘蛛磁力 搜索引擎大全，如何使用蜘蛛磁力查找磁力链接07TRAE 规则（Rules）配置指南：个人习惯、团队规范与最佳实践08NVIDIA显卡驱动、CUDA、cuDNN 和 TensorRT 版本匹配指南09阿里开源首个图像生成基础模型——Qwen-Image本地部署教程，超强中文渲染能力刷新SOTA！10TRAE Rules 实践：为项目配置 6A 工作流