pytorch-分类-检测-分割的dataset和dataloader创建

1.前言

在PyTorch中,DatasetDataLoader是两个重要的工具,用于构建输入数据的管道。

(1)Dataset是一个抽象类,表示数据集,需要实现__len____getitem__方法。

(2)DataLoader是一个可迭代的数据加载器,它封装了数据集的加载、批处理、打乱和并行加载等功能。

2.分类任务创建DatasetDataLoader

(1)对于分类任务,Dataset需要返回图像和对应的标签

python 复制代码
from torch.utils.data import Dataset  
from PIL import Image  
import os  
import torch  
  
class ClassificationDataset(Dataset):  
    def __init__(self, root_dir, transform=None):  
        self.transform = transform  
        self.images = [os.path.join(root_dir, img) for img in os.listdir(root_dir) if img.endswith('.jpg')]  
        self.labels = [...]  # 这里应该是与图像对应的标签列表  
  
    def __len__(self):  
        return len(self.images)  
  
    def __getitem__(self, idx):  
        img_path = self.images[idx]  
        image = Image.open(img_path).convert('RGB')  
        label = self.labels[idx]  
          
        if self.transform:  
            image = self.transform(image)  
          
        return image, label

(2)DataLoader加载数据

python 复制代码
from torch.utils.data import DataLoader  
  
transform = ...  # 这里定义你的数据预处理流程  
dataset = ClassificationDataset(root_dir='path_to_your_data', transform=transform)  
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

3.检测任务创建DatasetDataLoader

(1)Dataset需要返回图像和对应的边界框信息

python 复制代码
class DetectionDataset(Dataset):  
    def __init__(self, root_dir, transform=None):  
        self.transform = transform  
        self.images = [os.path.join(root_dir, img) for img in os.listdir(root_dir) if img.endswith('.jpg')]  
        self.annotations = [...]  # 这里应该是与图像对应的边界框信息列表  
  
    def __len__(self):  
        return len(self.images)  
  
    def __getitem__(self, idx):  
        img_path = self.images[idx]  
        image = Image.open(img_path).convert('RGB')  
        boxes = self.annotations[idx]  # 这些是边界框信息  
  
        if self.transform:  
            image, boxes = self.transform(image, boxes)  
          
        return image, boxes

(2)DataLoader加载数据

python 复制代码
dataloader = DataLoader(DetectionDataset(root_dir='path_to_your_data', transform=transform), batch_size=2, shuffle=True)

4.分割任务创建DatasetDataLoader

(1)Dataset需要返回图像和对应的分割掩码

python 复制代码
class SegmentationDataset(Dataset):  
    def __init__(self, root_dir, transform=None):  
        self.transform = transform  
        self.images = [os.path.join(root_dir, img) for img in os.listdir(root_dir) if img.endswith('.jpg')]  
        self.masks = [...]  # 这里应该是与图像对应的分割掩码列表  
  
    def __len__(self):  
        return len(self.images)  
  
    def __getitem__(self, idx):  
        img_path = self.images[idx]  
        mask_path = self.masks[idx]  
        image = Image.open(img_path).convert('RGB')  
        mask = Image.open(mask_path).convert('L')  # 假设掩码是灰度图  
  
        if self.transform:  
            image, mask = self.transform(image, mask)  
          
        return image, mask

(2)DataLoader加载数据

python 复制代码
dataloader = DataLoader(SegmentationDataset(root_dir='path_to_your_data', transform=transform), batch_size=4, shuffle=True)

在PyTorch的DatasetDataLoader框架中,idx(或称为索引)是通过迭代DataLoader时自动生成的。当你创建一个DataLoader实例,并在训练循环中迭代它时,DataLoader会内部调用Dataset__getitem__方法,并自动为你提供索引idx

相关推荐
X.AI66618 分钟前
YouTube评论情感分析项目84%正确率:基于BERT的实战复现与原理解析
人工智能·深度学习·bert
艾莉丝努力练剑24 分钟前
【C++:继承】面向对象编程精要:C++继承机制深度解析与最佳实践
开发语言·c++·人工智能·继承·c++进阶
小宁爱Python39 分钟前
从零搭建 RAG 智能问答系统 6:Text2SQL 与工作流实现数据库查询
数据库·人工智能·python·django
Hard_Liquor40 分钟前
Datawhale秋训营-“大运河杯”数据开发应用创新大赛
人工智能·深度学习·算法
运维行者_1 小时前
AWS云服务故障复盘——从故障中汲取的 IT 运维经验
大数据·linux·运维·服务器·人工智能·云计算·aws
Saniffer_SH1 小时前
搭载高性能GPU的英伟达Nvidia DGX Spark桌面性能小怪兽国内首台开箱视频!
人工智能·深度学习·神经网络·ubuntu·机器学习·语言模型·边缘计算
数字化脑洞实验室1 小时前
AI决策vs人工决策:效率的底层逻辑与选择边界
人工智能
可触的未来,发芽的智生1 小时前
追根索源:换不同的词嵌入(词向量生成方式不同,但词与词关系接近),会出现什么结果?
javascript·人工智能·python·神经网络·自然语言处理
递归不收敛1 小时前
三、检索增强生成(RAG)技术体系
人工智能·笔记·自然语言处理
kyle-fang2 小时前
pytorch-损失函数
pytorch