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

相关推荐
咨询187150651274 分钟前
高企复审奖补!2025年合肥市高新技术企业重新认定奖励补贴政策及申报条件
大数据·人工智能·区块链
Guheyunyi16 分钟前
智能照明系统:照亮智慧生活的多重价值
大数据·前端·人工智能·物联网·信息可视化·生活
云天徽上34 分钟前
【数据可视化-27】全球网络安全威胁数据可视化分析(2015-2024)
人工智能·安全·web安全·机器学习·信息可视化·数据分析
ONEYAC唯样40 分钟前
“在中国,为中国” 英飞凌汽车业务正式发布中国本土化战略
大数据·人工智能
mozun20201 小时前
产业观察:哈工大机器人公司2025.4.22
大数据·人工智能·机器人·创业创新·哈尔滨·名校
-一杯为品-1 小时前
【深度学习】#9 现代循环神经网络
人工智能·rnn·深度学习
硅谷秋水1 小时前
ORION:通过视觉-语言指令动作生成的一个整体端到端自动驾驶框架
人工智能·深度学习·机器学习·计算机视觉·语言模型·自动驾驶
Java中文社群1 小时前
最火向量数据库Milvus安装使用一条龙!
java·人工智能·后端
豆芽8191 小时前
强化学习(Reinforcement Learning, RL)和深度学习(Deep Learning, DL)
人工智能·深度学习·机器学习·强化学习
山北雨夜漫步1 小时前
机器学习 Day14 XGboost(极端梯度提升树)算法
人工智能·算法·机器学习