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

相关推荐
2301_821919927 分钟前
深度学习(四)
pytorch·深度学习
从孑开始14 分钟前
ManySpeech.MoonshineAsr 使用指南
人工智能·ai·c#·.net·私有化部署·语音识别·onnx·asr·moonshine
涛涛讲AI18 分钟前
一段音频多段字幕,让音频能够流畅自然对应字幕 AI生成视频,扣子生成剪映视频草稿
人工智能·音视频·语音识别
可触的未来,发芽的智生32 分钟前
新奇特:黑猫警长的纳米世界,忆阻器与神经网络的智慧
javascript·人工智能·python·神经网络·架构
WWZZ20251 小时前
快速上手大模型:机器学习2(一元线性回归、代价函数、梯度下降法)
人工智能·算法·机器学习·计算机视觉·机器人·大模型·slam
AKAMAI1 小时前
数据孤岛破局之战 :跨业务分析的难题攻坚
运维·人工智能·云计算
Chicheng_MA1 小时前
算能 CV184 智能相机整体方案介绍
人工智能·数码相机·算能
Element_南笙1 小时前
吴恩达新课程:Agentic AI(笔记2)
数据库·人工智能·笔记·python·深度学习·ui·自然语言处理
倔强青铜三2 小时前
苦练Python第69天:subprocess模块从入门到上瘾,手把手教你驯服系统命令!
人工智能·python·面试
Antonio9152 小时前
【图像处理】rgb和srgb
图像处理·人工智能·数码相机