当前位置：首页 > news >正文

从UFLD到UFLDv2实战：在自定义数据集上快速实现车道线检测（PyTorch版）

news 2026/5/5 5:43:29

从UFLD到UFLDv2实战：在自定义数据集上快速实现车道线检测（PyTorch版）

车道线检测是自动驾驶和机器人导航中的基础任务，而UFLD系列模型以其高效和准确的特点成为该领域的热门选择。本文将带您从零开始，在PyTorch框架下实现UFLD和UFLDv2模型，并应用于自定义数据集。

1. 环境配置与数据准备

1.1 基础环境搭建

首先需要配置PyTorch环境。推荐使用Python 3.8+和PyTorch 1.10+版本：

conda create -n ufld python=3.8 conda activate ufld pip install torch torchvision torchaudio pip install opencv-python pandas tqdm

对于GPU加速，确保安装对应CUDA版本的PyTorch。可以通过nvidia-smi查看CUDA版本。

1.2 数据集格式处理

UFLD系列模型通常使用CULane或TuSimple格式的数据集。自定义数据集需要转换为以下结构：

dataset/ ├── images/ │ ├── train/ │ │ ├── 0001.jpg │ │ └── ... │ └── val/ │ ├── 0001.jpg │ └── ... └── labels/ ├── train/ │ ├── 0001.lines.txt │ └── ... └── val/ ├── 0001.lines.txt └── ...

每个.lines.txt文件包含多行，每行表示一条车道线的坐标，格式为：

x1 y1 x2 y2 ... xn yn

提示：可以使用OpenCV的cv2.polylines函数可视化标注，确保数据标注正确。

2. UFLD模型实现

2.1 模型架构解析

UFLD的核心创新在于将车道检测转化为基于行锚的分类问题。其网络结构主要包含：

骨干网络：通常使用ResNet或EfficientNet提取特征
分类头：预测每个行锚点上车道的位置概率分布
结构损失：包括相似度损失和形状损失

import torch import torch.nn as nn class UFLD(nn.Module): def __init__(self, backbone='resnet18', num_lanes=4, num_anchors=72): super().__init__() # 骨干网络 self.backbone = torch.hub.load('pytorch/vision', backbone, pretrained=True) in_features = self.backbone.fc.in_features self.backbone = nn.Sequential(*list(self.backbone.children())[:-2]) # 分类头 self.cls_head = nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(in_features, num_lanes * num_anchors * (num_cells + 1)) ) def forward(self, x): features = self.backbone(x) logits = self.cls_head(features) return logits.view(-1, self.num_lanes, self.num_anchors, self.num_cells + 1)

2.2 损失函数实现

UFLD使用三种损失函数的组合：

分类损失：交叉熵损失
相似度损失：相邻行锚预测的L1距离
形状损失：二阶差分约束

def ufld_loss(pred, target): # 分类损失 cls_loss = F.cross_entropy(pred, target) # 相似度损失 pred_prob = F.softmax(pred[:, :, :-1], dim=-1) # 排除背景类 sim_loss = torch.mean(torch.abs(pred_prob[:, :, 1:] - pred_prob[:, :, :-1])) # 形状损失 loc = torch.sum(pred_prob * torch.arange(pred_prob.size(-1), device=pred_prob.device), dim=-1) shp_loss = torch.mean(torch.abs( (loc[:, :, 2:] - loc[:, :, 1:-1]) - (loc[:, :, 1:-1] - loc[:, :, :-2]) )) return cls_loss + 0.5 * sim_loss + 0.5 * shp_loss

3. UFLDv2改进与实现

3.1 混合锚点系统

UFLDv2的核心改进是引入了混合锚点系统：

特性	UFLD	UFLDv2
锚点类型	仅行锚	行锚+列锚
适用场景	垂直车道	所有方向车道
定位误差	水平车道误差大	各方向误差均衡
计算成本	低	中等

实现混合锚点需要修改网络结构：

class UFLDv2(nn.Module): def __init__(self, backbone='resnet34', num_row_anchors=72, num_col_anchors=40): super().__init__() self.backbone = torch.hub.load('pytorch/vision', backbone, pretrained=True) in_features = self.backbone.fc.in_features self.backbone = nn.Sequential(*list(self.backbone.children())[:-2]) # 行锚分支 self.row_head = nn.Sequential( nn.Conv2d(in_features, 256, kernel_size=1), nn.Flatten(), nn.Linear(256 * 8 * 8, num_row_anchors * (num_cells + 1)) ) # 列锚分支 self.col_head = nn.Sequential( nn.Conv2d(in_features, 256, kernel_size=1), nn.Flatten(), nn.Linear(256 * 8 * 8, num_col_anchors * (num_cells + 1)) )

3.2 有序分类损失

UFLDv2引入了有序分类的概念：

基础分类损失：标准交叉熵损失
期望损失：约束预测分布的期望接近真实值

def ufldv2_loss(row_pred, col_pred, row_target, col_target): # 基础分类损失 row_cls_loss = F.cross_entropy(row_pred, row_target) col_cls_loss = F.cross_entropy(col_pred, col_target) # 期望损失 row_prob = F.softmax(row_pred, dim=-1) row_exp = torch.sum(row_prob * torch.arange(row_prob.size(-1), device=row_prob.device), dim=-1) row_exp_loss = F.smooth_l1_loss(row_exp, row_target.float()) col_prob = F.softmax(col_pred, dim=-1) col_exp = torch.sum(col_prob * torch.arange(col_prob.size(-1), device=col_prob.device), dim=-1) col_exp_loss = F.smooth_l1_loss(col_exp, col_target.float()) return row_cls_loss + col_cls_loss + 0.3 * (row_exp_loss + col_exp_loss)

4. 训练与评估

4.1 训练流程优化

训练时需要注意以下关键点：

学习率调度：使用余弦退火学习率
数据增强：
- 随机水平翻转
- 颜色抖动
- 透视变换
批量大小：根据GPU内存选择最大可能值

from torch.optim.lr_scheduler import CosineAnnealingLR model = UFLDv2().cuda() optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) scheduler = CosineAnnealingLR(optimizer, T_max=100) for epoch in range(100): for images, targets in train_loader: images = images.cuda() row_targets, col_targets = targets # 前向传播 row_pred, col_pred = model(images) # 计算损失 loss = ufldv2_loss(row_pred, col_pred, row_targets, col_targets) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step()

4.2 评估指标与可视化

常用评估指标包括：

准确率：预测正确的车道点比例
FP/FN：误检/漏检率
F1分数：综合评估指标

可视化可以使用以下代码：

def visualize(image, predictions): image = image.copy() h, w = image.shape[:2] # 绘制行锚预测 for lane in predictions['row']: points = [(int(x * w), int(y * h)) for x, y in lane] cv2.polylines(image, [np.array(points)], False, (0, 255, 0), 2) # 绘制列锚预测 for lane in predictions['col']: points = [(int(x * w), int(y * h)) for x, y in lane] cv2.polylines(image, [np.array(points)], False, (255, 0, 0), 2) return image

5. 实际应用中的优化技巧

5.1 模型轻量化

对于嵌入式设备部署，可以考虑：

知识蒸馏：用大模型训练小模型
量化：FP16或INT8量化
剪枝：移除不重要的通道

# FP16混合精度训练示例 from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() with autocast(): predictions = model(images) loss = criterion(predictions, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()