BDD100K目标检测实战：从数据准备到模型训练的全流程指南

1次阅读

共计 3417 个字符，预计需要花费 9 分钟才能阅读完成。

BDD100K 是 UC Berkeley 发布的自动驾驶场景数据集，包含 10 万张高清图像，涵盖多样化的天气条件（晴天、雨天、雾天等）、光照变化（白天、夜晚）和复杂道路环境（城市、乡村、高速公路）。其标注包含 2D 边界框、可行驶区域、车道线等多任务标签。对于目标检测任务，主要挑战在于：

标注复杂性 ：平均每张图包含 3.4 个物体，但分布极不均匀（如车辆类占 70% 以上）
场景多样性 ：同一物体的外观差异大（如夜间车灯开启的车辆 vs 白天正常车辆）

数据加载瓶颈 ：直接使用原生 JSON 标注会导致 GPU 利用率不足 30%
类别失衡 ：行人类样本仅为车辆类的 1 /8，导致模型对少数类识别率低
预处理误区 ：盲目应用翻转增强可能破坏交通标志的可读性

推荐使用 Albumentations 库，相比传统 OpenCV 操作有 2 - 3 倍的加速：

import albumentations as A

transform = A.Compose([A.RandomBrightnessContrast(p=0.5),  # 亮度对比度调整
    A.HueSaturationValue(p=0.3),       # 色相饱和度调整
    A.RandomFog(fog_coef_lower=0.1, fog_coef_upper=0.3, p=0.1),  # 雾天模拟
    A.Cutout(max_h_size=32, max_w_size=32, p=0.2)  # 随机遮挡
], bbox_params=A.BboxParams(format='coco'))

采用提前解析 JSON 到二进制缓存的方法：

将原始 JSON 转换为 Parquet 格式，读取速度提升 5 倍
使用内存映射文件处理大尺寸图像
对类别 ID 进行哈希编码加速查找

模型	mAP@0.5	FPS (Tesla V100)	显存占用
Faster R-CNN	58.7	23	9.2GB
YOLOv5s	54.2	65	4.1GB
RetinaNet	56.1	38	6.8GB

选择建议 ：
– 追求精度：Faster R-CNN + ResNet101
– 需要实时性：YOLOv5s + 剪枝量化

PyTorch DDP 模式关键配置：

torch.distributed.init_process_group(
    backend='nccl',
    init_method='env://'
)
model = torch.nn.parallel.DistributedDataParallel(
    model,
    device_ids=[local_rank],
    output_device=local_rank
)

class BDD100KDataset(torch.utils.data.Dataset):
    def __init__(self, root, transform=None):
        self.annotations = pd.read_parquet('annotations.parquet')  # 预处理的标注
        self.image_dir = Path(root) / 'images'
        self.transform = transform

    def __getitem__(self, idx) -> Tuple[torch.Tensor, Dict]:
        record = self.annotations.iloc[idx]
        img_path = self.image_dir / record['image_name']

        # 使用内存映射加载大图
        img = np.load(img_path.with_suffix('.npy'), mmap_mode='r')

        # 转换为 Albumentations 需要的格式
        bboxes = record['bboxes'].reshape(-1, 4)  # [N,4]
        labels = record['class_ids']  # [N,]

        if self.transform:
            augmented = self.transform(image=img, bboxes=bboxes, class_labels=labels)
            img, bboxes, labels = augmented['image'], augmented['bboxes'], augmented['class_labels']

        # 转换为 Tensor
        target = {'boxes': torch.as_tensor(bboxes, dtype=torch.float32),
            'labels': torch.as_tensor(labels, dtype=torch.int64)
        }
        return torch.from_numpy(img).permute(2,0,1), target

class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, preds: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        ce_loss = F.cross_entropy(preds, targets, reduction='none')
        pt = torch.exp(-ce_loss)
        loss = self.alpha * (1-pt)**self.gamma * ce_loss
        return loss.mean()

梯度累积 ：每 4 个 batch 更新一次参数

for i, (images, targets) in enumerate(dataloader):
    preds = model(images)
    loss = criterion(preds, targets) / 4  # 梯度累积
    loss.backward()

    if (i+1) % 4 == 0:
        optimizer.step()
        optimizer.zero_grad()

混合精度训练 ：

scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
    preds = model(images)
    loss = criterion(preds, targets)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

关键转换参数：

trtexec --onnx=model.onnx \
        --fp16 \
        --workspace=4096 \
        --minShapes=input:1x3x640x640 \
        --optShapes=input:8x3x640x640 \
        --maxShapes=input:32x3x640x640

def validate_annotations(annotations):
    for ann in annotations:
        x1, y1, w, h = ann['bbox']
        assert w > 0 and h > 0, f"Invalid bbox size {ann['bbox']}"
        assert ann['category_id'] in VALID_CLASSES, f"Unknown class {ann['category_id']}"