如何为2D机器人地图矩阵数据创建PyTorch自定义Dataset数据集

核心逻辑 #

PyTorch的自定义Dataset不依赖特定的文件格式，只需要继承torch.utils.data.Dataset基类，实现__len__（返回数据集总样本数）和__getitem__（根据索引返回对应样本和标签）两个核心方法即可，完全适配你当前的多通道二值矩阵输入需求。
你的输入格式和三通道RGB图像逻辑完全一致：3个通道分别对应free、unknown、obstacle三类的二值掩码，形状为[3, 地图高度, 地图宽度]，可以直接输入CNN训练。

实现步骤 #

第一步：整理数据存储 #

你可以根据自己的使用场景选择数据存储方式：

• 小数据集：直接把所有样本和对应的标签加载到内存，存为numpy数组或者PyTorch张量列表
• 大数据集：把每个样本的输入矩阵、标签矩阵单独存为.npy（numpy格式）或者.pt（PyTorch序列化格式）文件，提前整理好所有文件的路径列表

第二步：编写自定义Dataset类 #

提供两种常见场景的代码示例：

场景1：数据已全部加载到内存

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np

class MapDataset(Dataset):
    def __init__(self, input_data, label_data, transform=None):
        """
        参数说明：
        input_data: 所有输入样本的集合，形状为[总样本数, 3, H, W]，3个通道分别对应三类标签的二值矩阵
        label_data: 所有对应标签的集合，形状为[总样本数, H, W]，每个位置的值为1/2/3对应三类标注
        transform: 可选，自定义数据增强操作
        """
        self.input_data = input_data
        self.label_data = label_data
        self.transform = transform
    
    def __len__(self):
        return len(self.input_data)
    
    def __getitem__(self, idx):
        # 按索引取对应样本，转成float32类型张量
        input_map = torch.tensor(self.input_data[idx], dtype=torch.float32)
        # 标签转成long类型，适配交叉熵损失要求
        label_map = torch.tensor(self.label_data[idx], dtype=torch.long)
        
        if self.transform:
            input_map = self.transform(input_map)
        
        return input_map, label_map

场景2：数据存在本地npy文件中

class MapDatasetFromFiles(Dataset):
    def __init__(self, input_file_paths, label_file_paths, transform=None):
        """
        参数说明：
        input_file_paths: 所有输入样本npy文件的路径列表
        label_file_paths: 所有对应标签npy文件的路径列表，顺序和输入列表一一对应
        transform: 可选，自定义数据增强操作
        """
        self.input_paths = input_file_paths
        self.label_paths = label_file_paths
        self.transform = transform
    
    def __len__(self):
        return len(self.input_paths)
    
    def __getitem__(self, idx):
        # 读取npy文件
        input_map = np.load(self.input_paths[idx])
        label_map = np.load(self.label_paths[idx])
        
        # 转成张量
        input_map = torch.tensor(input_map, dtype=torch.float32)
        label_map = torch.tensor(label_map, dtype=torch.long)
        
        if self.transform:
            input_map = self.transform(input_map)
        
        return input_map, label_map

第三步：调用DataLoader加载数据 #

实例化Dataset后直接传入DataLoader即可正常使用：

# 示例：使用内存数据集
# 假设你已经准备好了输入和标签数组
total_samples = 1000
H, W = 64, 64
mock_input = np.random.randint(0, 2, size=(total_samples, 3, H, W)) # 模拟二值掩码输入
mock_label = np.random.randint(1, 4, size=(total_samples, H, W)) # 模拟1/2/3的标签

dataset = MapDataset(mock_input, mock_label)
dataloader = DataLoader(
    dataset,
    batch_size=16,
    shuffle=True,
    num_workers=4 # 按需设置加载进程数，Windows环境如果报错可以设为0
)

# 训练时遍历dataloader即可
for batch_input, batch_label in dataloader:
    # 送入CNN训练的逻辑
    pass

注意事项 #

• 输入的三个通道顺序要固定，比如第0层对应free、第1层对应unknown、第2层对应obstacle，训练和推理阶段必须保持一致
• 如果需要做数据增强（比如随机翻转、裁剪地图），可以直接把对应操作写在transform参数里，逻辑和图像处理的增强完全通用
• 若训练目标不是逐像素分割，只需修改__getitem__返回的标签格式即可，比如返回整张地图的全局属性值

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