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

pytorch模型的进阶训练和性能优化

news 2025/4/3 14:45:15

综合案例

将MNIST数据集保存成本地图片
读取本地图片进行训练
读取自己的数据集进行训练
用自己的模型进行训练
获得更多评价指标
提升模型性能的方法

MNIST转本地图片

import os
import torchvision
import torchvision.transforms as transforms# 下载MNIST数据集
transform = transforms.Compose([transforms.ToTensor()])
mnist_trainset = torchvision.datasets.MNIST(root='./data', train=True,download=True, transform=transform)
mnist_testset = torchvision.datasets.MNIST(root='./data', train=False,download=True, transform=transform)# 创建一个目录来保存图像（如果它还不存在）
os.makedirs('./mnist_images/train', exist_ok=True)
os.makedirs('./mnist_images/test', exist_ok=True)# 遍历数据集并保存图像
for idx, (image, label) in enumerate(mnist_trainset):# 创建类别文件夹（如果它还不存在）label_dir = os.path.join('./mnist_images/train', str(label))os.makedirs(label_dir, exist_ok=True)# 转换为PIL图像并保存pil_image = transforms.ToPILImage()(image)pil_image.save(os.path.join(label_dir, f'{idx}.jpg'))# 遍历数据集并保存图像
for idx, (image, label) in enumerate(mnist_testset):# 创建类别文件夹（如果它还不存在）label_dir = os.path.join('./mnist_images/test', str(label))os.makedirs(label_dir, exist_ok=True)# 转换为PIL图像并保存pil_image = transforms.ToPILImage()(image)pil_image.save(os.path.join(label_dir, f'{idx}.jpg'))# 打印完成消息
print("All images have been saved successfully.")

接下来我来讲解一下上述的代码，在我的视角看来应该要将的东西

transform = transforms.Compose([transforms.ToTensor()])

使用 torchvision.transforms 模块中的 Compose 和 ToTensor 方法来定义一个图像预处理的转换操作，主要用于将图像数据转换为 PyTorch 张量（Tensor），以便用于深度学习模型的训练或推理。

mnist_trainset=torchvision.datasets.MNIST(root='./data',train=True,download=True, transform=transform

自动下载MNIST数据集，然后将其转换为tensor格式

os.makedirs('./mnist_images/train', exist_ok=True)

使用 Python 的 os 模块中的 makedirs 函数来创建目录。具体来说，它的作用是创建一个目录路径 ./mnist_images/train，并且如果该目录已经存在，不会报错。

其中exist_ok=True，在目录已经存在的情况下，不会报错

label_dir = os.path.join('./mnist_images/train', str(label))

这一句代码的作用就是，把'./mnist_images/train'字符串和str(label)字符串拼接起来。

`pil_image = transforms.ToPILImage()(image)`

这行代码的作用是将输入的图像数据（通常是 PyTorch 张量或 NumPy 数组）转换为 PIL 图像对象。

读取本地图片进行训练

import os
import torch
from torch.utils.data import Dataset, DataLoader
import cv2 as cvclass MNISTDataset(Dataset):def __init__(self, root_dir):self.root_dir = root_dirself.file_list = []self.name_list = []self.id_list = []for root, dirs, files in os.walk(self.root_dir):if dirs:self.name_list = dirsfor file_i in files:file_i_full_path = os.path.join(root, file_i)file_class = os.path.split(file_i_full_path)[0].split('\\')[-1]self.id_list.append(self.name_list.index(file_class))self.file_list.append(file_i_full_path)def __len__(self):return len(self.file_list)def __getitem__(self, idx):img = self.file_list[idx]img = cv.imread(img, 0)img = cv.resize(img, dsize=(28, 28))img = torch.from_numpy(img).float()label = self.id_list[idx]# print(label)label = torch.tensor(label)return img, labelif __name__ == '__main__':my_dataset_train = MNISTDataset(r'mnist_images/train')my_dataloader_train = DataLoader(my_dataset_train, batch_size=10, shuffle=True)# 尝试读取训练集数据print("读取训练集数据")for x, y in my_dataloader_train:print(x.type(), x.shape, y)my_dataset_test = MNISTDataset(r'mnist_images/test')my_dataloader_test = DataLoader(my_dataset_test, batch_size=10, shuffle=False)# 尝试读取训练集数据print("读取测试集数据")for x, y in my_dataloader_test:print(x.shape, y)

在前面讲了，Dataset的三件套，__init__，__len__，__getitem__如果这三个魔法方法忘记了，可以回去看看以前的文章。