如何在PyTorch中创建类数据集关联图像与标签用于图像分类

你完全可以通过自定义PyTorch的Dataset子类来关联图像张量与标签，这是PyTorch处理自定义分类数据集的标准方式。需要注意的是，PyTorch默认要求图像张量为通道在前的格式（即[C, H, W]），而你的图像当前是通道在后的[H, W, C]，所以需要在数据集中做维度转换。

自定义数据集类实现 #

import torch
from torch.utils.data import Dataset

class TensorImageDataset(Dataset):
    def __init__(self, images_tensor, labels_tensor, transform=None):
        # 校验图像与标签的样本数是否匹配
        assert len(images_tensor) == len(labels_tensor), "图像和标签的样本数量不匹配"
        self.images = images_tensor
        self.labels = labels_tensor
        self.transform = transform
        # 确保标签为分类任务常用的长整型
        self.labels = self.labels.long()

    def __len__(self):
        # 返回数据集总样本数
        return len(self.images)

    def __getitem__(self, idx):
        # 获取单样本的图像与标签
        img = self.images[idx]
        label = self.labels[idx]

        # 将通道在后的[H,W,C]转为通道在前的[C,H,W]
        img = img.permute(2, 0, 1)
        # 若图像像素值为0-255范围，转为0-1的float类型（按需调整）
        img = img.float() / 255.0

        # 应用自定义数据变换（如增强操作）
        if self.transform:
            img = self.transform(img)

        return img, label

数据集与数据加载器使用示例 #

# 假设你的图像张量为images，标签张量为labels
dataset = TensorImageDataset(images, labels)

# 用DataLoader封装数据集，支持批量加载、打乱等操作
from torch.utils.data import DataLoader
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

# 测试数据加载
for batch_imgs, batch_labels in dataloader:
    print(batch_imgs.shape)  # 输出应为 [32, 3, 300, 300]
    print(batch_labels.shape) # 输出应为 [32]
    break

关键注意点 #

• 如果你的图像已经是归一化后的0-1范围，可去掉img = img.float() / 255.0这一步。
• transform参数可传入torchvision.transforms中的变换操作，用于数据增强，不需要则设为None。
• 若标签原本不是长整型，必须转换为long()类型，否则分类损失函数会报错。

内容的提问来源于stack exchange，提问作者Chiam Yuwei