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

2025-03-16 约 2880 字预计阅读 6 分钟

https://bing.ee123.net/img/rand?artid=146290658

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

前言

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

一、Pytorch

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

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

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

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

作用

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

参数

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

4. 模型定义 CustomModel

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. 训练循环

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. 模型导出

#保存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

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

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

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. 训练配置与执行

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. 模型导出

#保存权重
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 和静态图模型

四、常见问题

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

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

num_workers 设置多少合适？

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

验证集是否需要 shuffle=False？

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

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

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