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

相关推荐
飞哥数智坊9 小时前
GPT-5-Codex 发布,Codex 正在取代 Claude
人工智能·ai编程
倔强青铜三9 小时前
苦练Python第46天:文件写入与上下文管理器
人工智能·python·面试
虫无涯10 小时前
Dify Agent + AntV 实战:从 0 到 1 打造数据可视化解决方案
人工智能
Dm_dotnet12 小时前
公益站Agent Router注册送200刀额度竟然是真的
人工智能
算家计算12 小时前
7B参数拿下30个世界第一!Hunyuan-MT-7B本地部署教程:腾讯混元开源业界首个翻译集成模型
人工智能·开源
机器之心13 小时前
LLM开源2.0大洗牌:60个出局,39个上桌,AI Coding疯魔,TensorFlow已死
人工智能·openai
Juchecar14 小时前
交叉熵:深度学习中最常用的损失函数
人工智能
林木森ai14 小时前
爆款AI动物运动会视频,用Coze(扣子)一键搞定全流程(附保姆级拆解)
人工智能·aigc
聚客AI15 小时前
🙋‍♀️Transformer训练与推理全流程:从输入处理到输出生成
人工智能·算法·llm
BeerBear16 小时前
【保姆级教程-从0开始开发MCP服务器】一、MCP学习压根没有你想象得那么难!.md
人工智能·mcp