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

Pytorch使用Dataset加载数据

news 2025/7/5 7:08:21

1、前言：

在阅读之前，需要配置好对应pytorch版本。
对于一般学习，使用cpu版本的即可。参考教程点我
导入pytorch包，使用如下命令即可。

import torch   # 注意虽然叫pytorch，但是在引用时是引用torch

2、神经网络获取数据

神经网络获取数据主要用到Dataset和Dataloader两个方法
Dataset主要用于获取数据以及对应的真实label
Dataloader主要为后面的网络提供不同的数据形式
在torch.utils.data包内提供了DataSet类，可在Pytorch官网看到对应的描述

class Dataset(Generic[T_co]):r"""An abstract class representing a :class:`Dataset`.All datasets that represent a map from keys to data samples should subclassit. All subclasses should overwrite :meth:`__getitem__`, supporting fetching adata sample for a given key. Subclasses could also optionally overwrite:meth:`__len__`, which is expected to return the size of the dataset by many:class:`~torch.utils.data.Sampler` implementations and the default optionsof :class:`~torch.utils.data.DataLoader`. Subclasses could alsooptionally implement :meth:`__getitems__`, for speedup batched samplesloading. This method accepts list of indices of samples of batch and returnslist of samples... note:::class:`~torch.utils.data.DataLoader` by default constructs an indexsampler that yields integral indices.  To make it work with a map-styledataset with non-integral indices/keys, a custom sampler must be provided."""def __getitem__(self, index) -> T_co:raise NotImplementedError("Subclasses of Dataset should implement __getitem__.")# def __getitems__(self, indices: List) -> List[T_co]:# Not implemented to prevent false-positives in fetcher check in# torch.utils.data._utils.fetch._MapDatasetFetcherdef __add__(self, other: "Dataset[T_co]") -> "ConcatDataset[T_co]":return ConcatDataset([self, other])# No `def __len__(self)` default?# See NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ]# in pytorch/torch/utils/data/sampler.py

根据上述描述可知，Dataset是一个抽象类，用于表示数据集。你可以通过继承这个类并实现以下方法来自定义数据集：

__len__(self): 返回数据集的大小，即数据集中有多少个样本。
__getitem__(self, idx): 根据索引 idx 返回数据集中的一个样本和对应的标签。

3、案例

使用Dataset读取文件夹E:\Python_learning\Deep_learning\dataset\hymenoptera_data\train\ants下所有图片。并获取对应的label，该数据集的文件夹的名字为对应的标签，而文件夹内为对应的训练集的图片。

import os
from torch.utils.data import Dataset
from PIL import Image
from torch.utils.tensorboard import SummaryWriter
from torchvision import transformsclass MyDataset(Dataset):def __init__(self, root_path, label):self.root_path = root_pathself.label = labelself.img_path = os.path.join(root_path, label)  # 拼接路径print(f"图片路径: {self.img_path}")  # 打印路径以进行调试try:self.img_path_list = os.listdir(self.img_path)  # 列出文件夹中的文件print(f"图片列表: {self.img_path_list}")  # 打印图片列表以进行调试except PermissionError as e:print(f"权限错误: {e}")except FileNotFoundError as e:print(f"文件未找到错误: {e}")def __getitem__(self, index):img_index = self.img_path_list[index]img_path = os.path.join(self.img_path, img_index)try:img = Image.open(img_path)except Exception as e:print(f"读取图片时出错: {e}, 图片路径: {img_path}")raise elabel = self.labelreturn img, labeldef __len__(self):return len(self.img_path_list)# 实例化这个类
my_data = MyDataset(root_path=r'E:\Python_learning\Deep_learning\dataset\hymenoptera_data\train', label='ants')
writer = SummaryWriter('logs')
for i in range(my_data.__len__()):img, label = my_data[i]  # 依次获取对应的图片# 此处img为PIL Image, 使用transforms中的ToTensor方法转化为tensor格式writer.add_image(tag=label, img_tensor=transforms.ToTensor()(img), global_step=i)
writer.close()
print(f"当前文件夹下{i + 1}张图片已读取完毕，请在Tensorboard中查看")