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

Dataset与DataLoader、transform

news 2026/2/9 19:26:54

文章目录

1、Dataset
2、DataLoader
- 2.1 参数详解
- - 2.1.1 num_works
  - 2.1.2 pin_memory
  - 2.1.3 collate_fn
3、图像增强
4、重写transform

1、Dataset

在 PyTorch 中，如果要创建自定义的数据集（Dataset），通常会继承 torch.utils.data.Dataset 类。这是因为 PyTorch 提供了 Dataset 类来帮助管理数据集，并定义了一些必要的方法和属性，如 len 和 getitem 等。

在这里插入图片描述

import torch.utils.data as data
import torch
class MyDataSet(data.Dataset):def __init__(self):#创建tensor 含有20个数值 0-19self.data=torch.arange(0,20)def __getitem__(self, index):return self.data[index]def __len__(self):return len(self.data)if __name__ == '__main__':dataset=MyDataSet()print(len(dataset)) #20print(dataset[3]) #tensor(3)

2、DataLoader

if __name__ == '__main__':dataset=MyDataSet()print(len(dataset)) #20print(dataset[3]) #tensor(3)#定义读取规则 这里是打乱顺序 每次读取4个 所以20/4=5 Dataloader的长度是5Dataloader=data.DataLoader(dataset,shuffle=True,batch_size=4)print(len(Dataloader))for i in Dataloader:print(i)   #tensor([ 5, 10, 15, 16])# tensor([ 7, 18,  4,  0])# tensor([11,  3, 14,  1])# tensor([19,  8, 13, 12])# tensor([ 6,  2,  9, 17])

2.1 参数详解

![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/95124230fc354

2.1.1 num_works

就是看用几个GPU来同时计算一个GPU计算一个batch_SIZE
在这里插入图片描述

2.1.2 pin_memory

加快读取速度默认为false

2.1.3 collate_fn

即我们DataLoader 每迭代出一个batch_size 是一个什么样的格式。
分类任务 其实collate_fn 不用重写用官方自带的即可
get_item 首先根据索引找到每张图片及标签然后是一个列表形式，之后再通过collate_fn 将图片和标签重组在一起，images的shape为 (batch_size,C,H,W).labels的shape为(batch_size,1）
在这里插入图片描述
对于图像分割或者关键点检测 我们都需要重写collate_fn.因为target是一个字典，传统的无法进行读取。
，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch。

 def collate_fn(batch):imgs_tuple, targets_tuple = tuple(zip(*batch))imgs_tensor = torch.stack(imgs_tuple)return imgs_tensor, targets_tuple

3、图像增强

数据增强可以增加训练集的样本数量，缓解过拟合，并提高模型的泛化能力，从而有效提升算法的性能

图像预处理：
1、将图像转换成tensor 的数据格式
2、将图像的像素值范围由 0~255 转换为 0~1
3、(height, width, channel) =====>>>(channel, height, width)
4、归一化图像
归一化可以优化算法的收敛速度和性能，和 BN 层的作用差不多
归一化处理还可以消除不同图像之间的亮度和颜色差异，提高模型的鲁棒性。
在这里插入图片描述

4、重写transform

传统的transform接口只能接受一个参数即图像本身。
重写transform的目的：可以接受多个参数，可以保证对我们的图像和标注进行同步处理，比如图像分类任务，如果我们对图像及进行了预处理，比如进行了图像裁剪和缩放以及旋转等，其对应的标注框也应该做同步变换，否则就会出错，这时候就需要我们重写transform，对图像和标注做同步处理。

比如：

class Compose(object):"""组合多个transform函数"""def __init__(self, transforms):self.transforms = transformsdef __call__(self, image, target):for t in self.transforms:image, target = t(image, target)return image, target

归一化只对图像做处理就好

class Normalize(object):def __init__(self, mean=None, std=None):self.mean = meanself.std = stddef __call__(self, image, target):image = F.normalize(image, mean=self.mean, std=self.std)return image, target