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

PyTorch 数据加载实战：从 CSV 到图像的全流程解析

web 2025/7/16 7:19:13

一、PyTorch 数据加载的核心组件

1.1 Dataset 类的核心方法

1.2 DataLoader 的作用

二、加载 CSV 数据实战

2.1 自定义 CSV 数据集

2.2 使用 TensorDataset 快速加载

三、加载图像数据实战

3.1 自定义图像数据集

3.2 使用 ImageFolder 快速加载

四、加载官方数据集

五、总结

在深度学习项目中，数据加载是模型训练的第一步，也是至关重要的一步。PyTorch 提供了灵活的数据加载工具，让我们能够轻松处理各种类型的数据。本文将结合实际代码，详细讲解如何使用 PyTorch 加载 CSV 数据和图像数据，帮助初学者快速掌握数据加载的核心技巧。

一、PyTorch 数据加载的核心组件

PyTorch 的数据加载主要依赖两个核心类：Dataset和DataLoader。

Dataset：负责数据的读取和预处理，是所有自定义数据集的基类
DataLoader：负责批量加载数据，支持打乱顺序、多线程加载等功能

1.1 Dataset 类的核心方法

自定义数据集需要继承Dataset类，并实现以下三个方法：

class CustomDataset(Dataset):def __init__(self, ...):  # 初始化数据集，加载文件路径等passdef __len__(self):  # 返回数据集大小return len(self.data)def __getitem__(self, index):  # 根据索引返回样本return sample, label

1.2 DataLoader 的作用

DataLoader像是一个 "搬运工"，将Dataset中的数据按批次搬运给模型：

dataloader = DataLoader(dataset=dataset,  # 要加载的数据集batch_size=32,    # 批次大小shuffle=True,     # 是否打乱数据num_workers=2     # 多线程加载
)

二、加载 CSV 数据实战

CSV 文件是存储表格数据的常用格式，比如学生成绩表、特征数据表等。下面我们通过实际代码讲解如何加载 CSV 数据。

2.1 自定义 CSV 数据集

import torch
from torch.utils.data import Dataset, DataLoader
import pandas as pdclass CsvDataset(Dataset):def __init__(self, filepath):# 读取CSV文件df = pd.read_csv(filepath)# 删除不需要的列（学号、姓名）df.drop(['学号', '姓名'], axis=1, inplace=True)# 提取特征和标签x = df.iloc[1:, :-1]  # 从第二行开始，取除最后一列外的所有列作为特征y = df.iloc[1:, -1]   # 从第二行开始，取最后一列作为标签# 转换为Tensorself.data = torch.tensor(x.values, dtype=torch.float)self.labels = torch.tensor(y.values, dtype=torch.float)def __len__(self):return len(self.data)def __getitem__(self, index):sample = self.data[index]label = self.labels[index]return sample, label# 测试代码
def test_csv_dataset():filepath = '大数据答辩成绩表.csv'dataset = CsvDataset(filepath)print(f"数据集大小: {len(dataset)}")print(f"第一个样本: {dataset[0]}")test_csv_dataset()

2.2 使用 TensorDataset 快速加载

如果数据已经是 Tensor 格式，可以使用TensorDataset快速创建数据集，无需自定义类：

def test_tensor_dataset():filepath = '大数据答辩成绩表.csv'df = pd.read_csv(filepath)df.drop(['学号', '姓名'], axis=1, inplace=True)x = df.iloc[1:, :-1]y = df.iloc[1:, -1]# 转换为Tensordata = torch.tensor(x.values, dtype=torch.float)labels = torch.tensor(y.values, dtype=torch.float)# 使用TensorDatasetdataset = TensorDataset(data, labels)print(f"第一个样本: {dataset[0]}")

三、加载图像数据实战

处理图像数据时，我们需要考虑图像的读取、大小调整、格式转换等问题。下面介绍两种加载图像数据的方法。

3.1 自定义图像数据集

import os
import cv2
from torch.utils.data import Datasetclass PicDataset(Dataset):def __init__(self, filepath):self.filepaths = []  # 存储图像路径self.labels = []     # 存储标签dirnames = []        # 存储类别名称# 遍历文件夹for root, dirs, files in os.walk(filepath):if len(dirs) > 0:dirnames = dirs  # 获取类别文件夹名称for file in files:f_path = os.path.join(root, file)self.filepaths.append(f_path)# 根据文件夹名称确定标签classname = root.split('\\')[-1]self.labels.append(dirnames.index(classname))def __len__(self):return len(self.filepaths)def __getitem__(self, index):filepath = self.filepaths[index]# 读取图像img = cv2.imread(filepath)# 调整图像大小为112x112img = cv2.resize(img, (112, 112))# 转换为Tensor并调整维度 (HWC -> CHW)t_img = torch.tensor(img)t_img = t_img.permute(2, 0, 1)label = self.labels[index]return t_img, label# 测试代码
def test_pic_dataset():filepath = r'E:\人工智能\深度学习\dataset\butterfly'dataset = PicDataset(filepath)print(f"数据集大小: {len(dataset)}")img, label = dataset[0]print(f"图像形状: {img.shape}, 标签: {label}")

3.2 使用 ImageFolder 快速加载

PyTorch 的ImageFolder是加载图像数据集的便捷工具，特别适合以下结构的数据集：

root/class1/img1.jpgimg2.jpgclass2/img1.jpgimg2.jpg

使用方法如下：

from torchvision.datasets import ImageFolder
from torchvision import transformsdef test_image_folder():filepath = r'E:\人工智能\深度学习\dataset\butterfly'# 定义图像转换transform = transforms.Compose([transforms.Resize((112, 112)),  # 调整大小transforms.ToTensor(),          # 转换为Tensor])# 使用ImageFolder加载数据dataset = ImageFolder(root=filepath, transform=transform)print(f"类别: {dataset.classes}")print(f"数据集大小: {len(dataset)}")# 创建DataLoaderdataloader = DataLoader(dataset=dataset,batch_size=1,shuffle=True)# 显示一张图像for img, label in dataloader:print(f"图像形状: {img.shape}")print(f"标签: {label}")breaktest_image_folder()

四、加载官方数据集

PyTorch 提供了许多常用的公开数据集（如 MNIST、CIFAR 等），可以直接下载使用：

from torchvision import datasets, transformsdef test_mnist_dataset():# 定义转换transform = transforms.Compose([transforms.ToTensor()])# 加载MNIST训练集dataset = datasets.MNIST(root='../dataset',  # 数据保存路径train=True,         # 训练集download=True,      # 如果没有数据则下载transform=transform)# 创建DataLoaderdataloader = DataLoader(dataset=dataset,batch_size=1,shuffle=True)# 显示一张图像for img, label in dataloader:print(f"图像形状: {img.shape}")print(f"标签: {label}")breaktest_mnist_dataset()