网站建设前期目标企业注册查询网

---

---

前言

本文简单介绍了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_dirself.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 = transformdef __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.0for 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.0correct = 0with 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_dirself.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)),  # 调整图像尺寸为224x224transforms.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.0for images, labels in train_loader:images = images.to(device)  # 数据移至GPUlabels = 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_dirself.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 = transformdef __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, labeldef __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数组，无需转为PILtransforms.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 和静态图模型

四、常见问题

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

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

2. num_workers 设置多少合适？

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

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

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

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

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

---