3. 构建数据加载器

任务目标

本案例实现一个PyTorch数据加载器，用于端到端自动驾驶模仿学习。该加载器从磁盘高效批量读取前置摄像头图像及其对应的转向角、油门和刹车标签，并输出为模型训练可直接使用的张量批次。

数据集约定

假设数据按以下目录结构存放，每个样本包含一张图像和一个标签文件（JSON格式）：

dataset/
├── images/
│   ├── 00001.jpg
│   ├── 00002.jpg
│   └── ...
└── labels/
    ├── 00001.json
    ├── 00002.json
    └── ...

每个JSON文件内容示例：{"steer": 0.12, "throttle": 0.65, "brake": 0.0}

完整实现

以下代码实现了一个可配置的PyTorch Dataset与DataLoader组合，支持图像缩放、归一化、标签归一化，并使用多进程预取加速数据流水线。

import json
import os
from pathlib import Path

import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import torchvision.transforms as T

class DrivingDataset(Dataset):
    """端到端驾驶数据集：加载图像 + 转向/油门/刹车标签"""

    def __init__(self, root_dir, img_size=(160, 320), normalize_labels=True):
        self.img_dir = Path(root_dir) / "images"
        self.label_dir = Path(root_dir) / "labels"
        self.img_size = img_size
        self.normalize_labels = normalize_labels

        # 按文件名排序保证顺序一致
        self.img_files = sorted(self.img_dir.glob("*.jpg"))
        self.label_files = sorted(self.label_dir.glob("*.json"))
        assert len(self.img_files) == len(self.label_files), \
            f"图像({len(self.img_files)})与标签({len(self.label_files)})数量不匹配"

        # 图像变换：缩放 + 转张量 + 归一化(ImageNet统计)
        self.transform = T.Compose([
            T.Resize(img_size),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
        ])

    def __len__(self):
        return len(self.img_files)

    def __getitem__(self, idx):
        # 加载图像
        img_path = self.img_files[idx]
        image = Image.open(img_path).convert("RGB")
        image = self.transform(image)

        # 加载标签
        label_path = self.label_files[idx]
        with open(label_path, "r") as f:
            label_data = json.load(f)

        steer = torch.tensor(label_data["steer"], dtype=torch.float32)
        throttle = torch.tensor(label_data["throttle"], dtype=torch.float32)
        brake = torch.tensor(label_data["brake"], dtype=torch.float32)

        if self.normalize_labels:
            # 转向角归一化到[-1, 1]（假设原始范围约[-1, 1]）
            steer = steer.clamp(-1.0, 1.0)
            # 油门/刹车归一化到[0, 1]（假设原始范围[0, 1]）
            throttle = throttle.clamp(0.0, 1.0)
            brake = brake.clamp(0.0, 1.0)

        # 拼接为单一标签向量 [steer, throttle, brake]
        labels = torch.stack([steer, throttle, brake])
        return image, labels

def create_dataloader(root_dir, batch_size=32, shuffle=True,
                      num_workers=4, img_size=(160, 320)):
    """创建训练/验证数据加载器"""
    dataset = DrivingDataset(root_dir, img_size=img_size)
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=True,          # GPU训练时加速主机到设备传输
        drop_last=False
    )
    return loader

# 最小调用示例
if __name__ == "__main__":
    loader = create_dataloader("./dataset", batch_size=8, num_workers=2)
    images, labels = next(iter(loader))
    print(f"图像批次形状: {images.shape}")   # [8, 3, 160, 320]
    print(f"标签批次形状: {labels.shape}")   # [8, 3]
    print(f"标签范围: steer [{labels[:,0].min():.2f}, {labels[:,0].max():.2f}]")

运行与验证

依赖安装：

pip install torch torchvision pillow

准备最小测试数据：创建dataset/images/与dataset/labels/，放入至少2对同名.jpg/.json文件即可运行。

执行验证：

python driving_loader.py

通过条件：程序无异常退出，输出图像批次形状为[batch_size, 3, 160, 320]，标签批次形状为[batch_size, 3]，且标签数值在合理范围内（steer∈[-1,1]，throttle/brake∈[0,1]）。

关键说明

• 图像尺寸默认160×320（高×宽），可根据模型输入调整img_size参数。
• 标签归一化采用截断方式，不改变原始分布，仅保证数值范围有效。
• 若需加载灰度图或不同图像格式，修改glob("*.jpg")与convert("RGB")即可。
• 实际训练时建议将num_workers设为CPU核心数，并启用pin_memory=True以提升GPU传输效率。