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

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，替换为对应模型结构。

相关推荐
小狗爱吃黄桃罐头18 分钟前
串口自动化断电测试
运维·python·自动化
niuTaylor24 分钟前
自动驾驶工程师之多传感器融合篇
人工智能·机器学习·自动驾驶
你喜欢喝可乐吗？28 分钟前
在 Ubuntu 中用 Docker 安装 RAGFlow
运维·服务器·人工智能·ubuntu·docker
AcrelGHP30 分钟前
工厂能耗系统完整解决方案 ——安科瑞企业能源管控平台
人工智能·能源
yukai0800832 分钟前
【最后203篇系列】020 rocksdb agent
python
Code blocks32 分钟前
小试牛刀-Turbine数据分发
python·算法·区块链
liuweidong080241 分钟前
【Pandas】pandas Series plot.area
python·信息可视化·pandas
明明跟你说过1 小时前
【Transformer】架构：解锁自然语言处理的无限可能
人工智能·深度学习·机器学习·ai·transformer
老大白菜1 小时前
Python八字排盘系统实现分析
开发语言·python
钱彬 (Qian Bin)1 小时前
QT Quick（C++）跨平台应用程序项目实战教程 3 — 项目基本设置（窗体尺寸、中文标题、窗体图标、可执行程序图标）
c++·人工智能·音乐播放器·qml·界面设计·qt quick
热门推荐
01RAG 实践- Ollama+RagFlow 部署本地知识库02从零安装 LLaMA-Factory 微调 Qwen 大模型成功及所有的坑03DeepSeek各版本说明与优缺点分析04本地化部署AI知识库：基于Ollama+DeepSeek+AnythingLLM保姆级教程05如何在WPS和Word/Excel中直接使用DeepSeek功能06本地部署DeepSeek教程（Mac版本）073.18练习08KGG转MP3工具|非KGM文件|解密音频09DeepSeek R1本地化部署 Ollama + Chatbox 打造最强 AI 工具10MySQL（七）MySQL和Oracle、PostgreSQL的区别