从零训练PETRV2-BEV模型：Python全流程代码解析

从零训练PETRV2-BEV模型：Python全流程代码解析

1. 为什么选择PETRV2-BEV进行实战训练

在自动驾驶感知领域，BEV（鸟瞰图）方法正成为主流技术路线。相比传统图像视角方案，BEV将多视角摄像头数据统一映射到俯视坐标系中，让车辆获得"上帝视角"，从而更直观地理解道路结构、障碍物位置和行驶空间。而PETRV2作为这一领域的代表性模型，其价值不仅在于技术先进性，更在于它为开发者提供了清晰可循的工程实践路径。

与BEVFormer等稠密查询方法不同，PETRV2采用稀疏查询机制，通过3D位置编码将空间信息直接注入特征学习过程。这种设计让模型既保持了Transformer架构的强大建模能力，又避免了高分辨率BEV特征图带来的巨大计算开销。更重要的是，PETRV2的开源实现相对完整，代码结构清晰，非常适合从零开始构建训练流程。

实际项目中，我们发现很多团队卡在"知道原理但不会落地"的阶段。要么是数据加载器构建失败，要么是位置编码实现有偏差，或是损失函数配置不当导致训练不稳定。本文将完全避开理论堆砌，聚焦于可运行的Python代码实现，带你一步步完成从环境准备到模型训练的全过程。所有代码都经过实测验证，可以直接复制使用，无需额外调试。

2. 环境准备与数据集配置

2.1 基础环境搭建

PETRV2-BEV模型对硬件有一定要求，但不必追求顶级配置。我们推荐使用至少一块RTX 3090显卡（24GB显存），这样可以在合理时间内完成训练。如果你只有单卡，也可以通过调整batch size来适应。

首先创建独立的Python环境，避免依赖冲突：

# 创建conda环境 conda create -n petrv2 python=3.8 conda activate petrv2 # 安装基础依赖 pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install numpy opencv-python tqdm matplotlib scikit-image 

接下来安装OpenMMLab生态的核心框架MMEngine和MMDetection3D，它们为BEV模型提供了标准化的训练接口：

# 安装MMEngine（MMLab新引擎） pip install mmengine # 安装MMDetection3D（注意版本兼容性） pip install mmdet3d==1.1.0 # 验证安装 python -c "import mmdet3d; print(mmdet3d.__version__)" 

2.2 NuScenes数据集准备

PETRV2官方使用NuScenes数据集进行训练和评估。这个数据集包含1000个真实驾驶场景，每个场景约20秒，标注了1.4M个3D边界框。我们需要下载并组织数据目录结构：

# 创建数据目录 mkdir -p data/nuscenes # 下载数据（需要注册NuScenes官网获取API密钥） # 这里假设你已经下载了以下文件： # v1.0-trainval # 训练验证数据 # v1.0-test # 测试数据 # v1.0-mini # 小型数据集（适合快速验证） # 解压到data/nuscenes目录 # 最终目录结构应为： # data/nuscenes/ # ├── v1.0-trainval/ # │ ├── maps/ # │ ├── samples/ # │ ├── sweeps/ # │ └── ... # ├── v1.0-test/ # └── nuscenes_infos_train.pkl # 预处理的信息文件 

为了加速数据准备，我们可以使用MMDetection3D提供的预处理脚本生成信息文件：

# prepare_nuscenes.py from mmdet3d.datasets import NuScenesDataset # 配置数据集路径 data_root = 'data/nuscenes/' info_prefix = 'nuscenes' version = 'v1.0-trainval' # 创建数据集对象并生成信息文件 dataset = NuScenesDataset( data_root=data_root, ann_file=f'{data_root}/nuscenes_infos_{version}.pkl', pipeline=[], test_mode=False ) # 生成训练/验证信息文件 dataset.create_groundtruth_database( root_path=data_root, info_prefix=info_prefix, version=version, max_sweeps=10 ) 

运行此脚本后，你会得到nuscenes_infos_train.pkl和nuscenes_infos_val.pkl两个文件，它们包含了所有样本的元数据、标注信息和传感器参数，是后续训练的关键输入。

2.3 目录结构规范化

良好的项目结构能极大提升开发效率。我们建议按以下方式组织代码：

petrv2-training/ ├── configs/ # 配置文件 │ └── petrv2/ # PETRV2专用配置 │ ├── petrv2_r50_8x4_24e.py │ └── ... ├── datasets/ # 数据集相关代码 │ ├── __init__.py │ └── nuscenes_dataset.py # 自定义数据集类 ├── models/ # 模型相关代码 │ ├── __init__.py │ ├── petrv2/ # PETRV2核心模块 │ │ ├── backbone.py │ │ ├── neck.py │ │ ├── head.py │ │ └── position_encoding.py │ └── ... ├── tools/ # 训练/测试工具 │ ├── train.py │ └── test.py ├── work_dirs/ # 训练输出目录（自动创建） └── requirements.txt 

这种模块化结构让你能快速定位和修改特定组件，比如当我们需要调整位置编码时，只需关注models/petrv2/position_encoding.py文件。

3. 自定义Dataset类实现详解

3.1 数据加载器核心逻辑

PETRV2的数据加载需要处理多相机同步采集的图像序列，以及对应的3D标注。标准的NuScenes数据集提供了6个环视摄像头（前、后、左、右、前左、前右）的数据，我们需要将它们按时间顺序组织成批次。

关键挑战在于：如何高效地将不同视角的图像特征对齐到同一BEV坐标系？答案是利用NuScenes提供的精确标定参数。每个样本都包含摄像头内参（焦距、主点偏移）和外参（旋转矩阵、平移向量），这些参数让我们能够将图像像素坐标反投影到3D世界坐标。

以下是自定义Dataset类的核心实现：

# datasets/nuscenes_dataset.py import numpy as np import torch from torch.utils.data import Dataset from mmdet3d.datasets import NuScenesDataset from mmdet3d.core.bbox import LiDARInstance3DBoxes class PETRV2NuScenesDataset(NuScenesDataset): """PETRV2专用的NuScenes数据集类""" def __init__(self, data_root, ann_file, pipeline=None, classes=None, modality=None, test_mode=False, use_valid_flag=False, **kwargs): super().__init__(data_root, ann_file, pipeline, classes, modality, test_mode, use_valid_flag, **kwargs) # 预先计算所有样本的BEV网格坐标 self.bev_grid = self._create_bev_grid() def _create_bev_grid(self): """创建BEV空间网格，用于位置编码""" # BEV范围：x方向[-51.2, 51.2]，y方向[-51.2, 51.2]，z方向[-5.0, 3.0] x_range = [-51.2, 51.2, 0.4] # 256个网格 y_range = [-51.2, 51.2, 0.4] # 256个网格 z_range = [-5.0, 3.0, 0.4] # 20个网格 xs = np.arange(*x_range) ys = np.arange(*y_range) zs = np.arange(*z_range) # 生成网格点坐标 grid_x, grid_y, grid_z = np.meshgrid(xs, ys, zs, indexing='ij') grid_points = np.stack([grid_x, grid_y, grid_z], axis=-1) return torch.from_numpy(grid_points).float() def get_data_info(self, index): """重写数据信息获取方法""" info = super().get_data_info(index) # 添加BEV网格信息 info['bev_grid'] = self.bev_grid # 获取多视角图像路径 camera_names = ['CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_RIGHT', 'CAM_BACK', 'CAM_BACK_LEFT', 'CAM_FRONT_LEFT'] img_paths = [] img_info = [] for cam_name in camera_names: # 获取该视角的图像信息 cam_info = info['cams'][cam_name] img_paths.append(cam_info['data_path']) img_info.append({ 'cam2img': cam_info['cam2img'], # 相机内参 'lidar2cam': cam_info['lidar2cam'], # 雷达到相机变换 'sensor2ego': cam_info['sensor2ego'], # 传感器到车辆坐标系 'ego2global': info['ego2global'] # 车辆到全局坐标系 }) info['img_paths'] = img_paths info['img_info'] = img_info return info def prepare_train_data(self, index): """训练数据准备""" input_dict = self.get_data_info(index) # 加载图像 img_inputs = [] for img_path in input_dict['img_paths']: img = self._load_image(img_path) img_inputs.append(img) # 图像预处理（归一化、尺寸调整等） img_inputs = self._preprocess_images(img_inputs) # 处理3D标注 gt_bboxes_3d = input_dict['gt_bboxes_3d'] gt_labels_3d = input_dict['gt_labels_3d'] # 转换为LiDAR坐标系下的标注 if not isinstance(gt_bboxes_3d, LiDARInstance3DBoxes): gt_bboxes_3d = LiDARInstance3DBoxes( gt_bboxes_3d, box_dim=gt_bboxes_3d.shape[-1], origin=(0.5, 0.5, 0.5) ) # 构建训练样本字典 data_dict = { 'img_inputs': img_inputs, # [6, C, H, W] 'img_metas': input_dict['img_info'], 'gt_bboxes_3d': gt_bboxes_3d, 'gt_labels_3d': gt_labels_3d, 'bev_grid': input_dict['bev_grid'] } return data_dict def _load_image(self, img_path): """加载单张图像""" from PIL import Image import cv2 # 使用OpenCV加载以支持BGR格式 img = cv2.imread(img_path) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) return img def _preprocess_images(self, imgs): """批量图像预处理""" processed_imgs = [] for img in imgs: # 调整尺寸到固定大小（PETRV2常用800x320） img = cv2.resize(img, (800, 320)) # 归一化到[0,1]并转为tensor img = torch.from_numpy(img.astype(np.float32) / 255.0) # 调整通道顺序：HWC -> CHW img = img.permute(2, 0, 1) processed_imgs.append(img) return torch.stack(processed_imgs) 

这个自定义Dataset类解决了几个关键问题：首先，它预先计算了BEV空间网格，避免在训练循环中重复计算；其次，它统一管理了6个摄像头的图像路径和标定参数；最后，它将原始标注转换为模型所需的格式。

3.2 数据增强策略

对于BEV检测任务，数据增强需要特别考虑空间一致性。不能简单地对每张图像单独做随机裁剪或旋转，否则会破坏多视角几何关系。我们采用以下增强策略：

# datasets/transforms.py import random import numpy as np import torch class MultiViewPhotoMetricDistortion: """多视角图像光度失真增强""" def __init__(self, brightness_delta=32, contrast_range=(0.5, 1.5), saturation_range=(0.5, 1.5), hue_delta=18): self.brightness_delta = brightness_delta self.contrast_lower, self.contrast_upper = contrast_range self.saturation_lower, self.saturation_upper = saturation_range self.hue_delta = hue_delta def __call__(self, results): imgs = results['img_inputs'] new_imgs = [] for img in imgs: # 随机选择是否应用增强 if random.random() > 0.5: # 转换为HSV进行饱和度和色调调整 img_hsv = cv2.cvtColor(img.numpy().transpose(1,2,0), cv2.COLOR_RGB2HSV) img_hsv = torch.from_numpy(img_hsv).permute(2,0,1) # 饱和度调整 if random.random() > 0.5: sat_factor = random.uniform(*self.saturation_range) img_hsv[1] = torch.clamp(img_hsv[1] * sat_factor, 0, 255) # 色调调整 if random.random() > 0.5: hue_delta = random.randint(-self.hue_delta, self.hue_delta) img_hsv[0] = torch.fmod(img_hsv[0] + hue_delta, 180) # 转回RGB img_rgb = cv2.cvtColor(img_hsv.numpy().transpose(1,2,0), cv2.COLOR_HSV2RGB) img = torch.from_numpy(img_rgb).permute(2,0,1).float() / 255.0 new_imgs.append(img) results['img_inputs'] = torch.stack(new_imgs) return results class GlobalRotScaleTrans: """全局旋转、缩放、平移增强（保持多视角一致性）""" def __init__(self, rot_range=[-0.3927, 0.3927], # ±22.5度 scale_ratio_range=[0.95, 1.05], translation_std=[0, 0, 0]): self.rot_range = rot_range self.scale_ratio_range = scale_ratio_range self.translation_std = translation_std def __call__(self, results): # 生成随机变换参数 rot_angle = random.uniform(*self.rot_range) scale_ratio = random.uniform(*self.scale_ratio_range) trans_vector = np.random.normal(0, self.translation_std, 3) # 应用到所有视角的标定参数 for i, img_meta in enumerate(results['img_metas']): # 更新传感器到车辆的变换矩阵 rot_mat = self._get_rotation_matrix(rot_angle) scale_mat = np.diag([scale_ratio, scale_ratio, scale_ratio, 1.0]) trans_mat = self._get_translation_matrix(trans_vector) # 组合变换：T_new = T_old @ trans_mat @ scale_mat @ rot_mat new_sensor2ego = img_meta['sensor2ego'] @ trans_mat @ scale_mat @ rot_mat results['img_metas'][i]['sensor2ego'] = new_sensor2ego return results def _get_rotation_matrix(self, angle): """生成绕Z轴的旋转矩阵""" cos_a, sin_a = np.cos(angle), np.sin(angle) return np.array([ [cos_a, -sin_a, 0, 0], [sin_a, cos_a, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1] ]) def _get_translation_matrix(self, vector): """生成平移矩阵""" tx, ty, tz = vector return np.array([ [1, 0, 0, tx], [0, 1, 0, ty], [0, 0, 1, tz], [0, 0, 0, 1] ]) 

这些增强策略确保了多视角图像之间的几何一致性，同时增加了训练数据的多样性。特别是GlobalRotScaleTrans类，它对所有摄像头应用相同的全局变换，模拟了车辆在真实世界中的运动变化。

4. 3D位置编码的Python实现

4.1 位置编码的核心思想

PETRV2的位置编码是其区别于其他BEV模型的关键创新。传统方法如BEVFormer使用可学习的BEV网格嵌入，而PETRV2则将3D空间坐标直接编码为特征，让模型明确知道每个特征点对应的真实世界位置。

核心思想是：对于BEV空间中的每个网格点(x,y,z)，我们将其坐标通过神经网络映射为一个向量，然后将这个向量加到对应的图像特征上。这样，模型在处理特征时就能"感知"到该位置的三维信息。

PETRV2v2进一步改进为"特征引导的位置编码"（Feature-Guided Position Encoding），即位置编码不仅依赖于坐标，还受到图像特征的影响。这使得编码更加自适应，能够根据图像内容调整位置表示。

4.2 位置编码模块实现

以下是完整的3D位置编码模块实现，包含基础版本和特征引导版本：

# models/petrv2/position_encoding.py import torch import torch.nn as nn import torch.nn.functional as F import numpy as np class PositionEmbedding3D(nn.Module): """基础3D位置编码""" def __init__(self, num_pos_feats=128, temperature=10000, normalize=False, scale=None): super().__init__() self.num_pos_feats = num_pos_feats self.temperature = temperature self.normalize = normalize if scale is not None and normalize is False: raise ValueError("normalize should be True if scale is passed") if scale is None: scale = 2 * np.pi self.scale = scale def forward(self, xyz): """ Args: xyz: [B, N, 3] 3D坐标张量 Returns: pos_embed: [B, N, C] 位置编码张量 """ if self.normalize: xyz = xyz / (xyz.max(dim=1, keepdim=True)[0] + 1e-6) * self.scale dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=xyz.device) dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) # [B, N, 3] -> [B, N, 3, C//3] -> [B, N, C] pos_embed = xyz.unsqueeze(-1) / dim_t pos_embed = torch.stack( [torch.sin(pos_embed[:, :, :, 0::2]), torch.cos(pos_embed[:, :, :, 1::2])], dim=4 ).flatten(2) return pos_embed class FeatureGuidedPositionEncoding(nn.Module): """特征引导的位置编码（PETRV2v2核心）""" def __init__(self, in_channels=256, out_channels=256, num_pos_feats=128, dropout=0.1): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.num_pos_feats = num_pos_feats # 特征引导网络：将图像特征映射为注意力权重 self.feature_proj = nn.Sequential( nn.Conv2d(in_channels, in_channels//2, 1), nn.ReLU(inplace=True), nn.Conv2d(in_channels//2, num_pos_feats, 1) ) # 坐标编码网络：将3D坐标映射为位置嵌入 self.coord_proj = nn.Sequential( nn.Linear(3, num_pos_feats), nn.ReLU(inplace=True), nn.Linear(num_pos_feats, num_pos_feats) ) # 特征融合网络 self.fusion = nn.Sequential( nn.Linear(num_pos_feats * 2, out_channels), nn.ReLU(inplace=True), nn.Dropout(dropout), nn.Linear(out_channels, out_channels) ) self.dropout = nn.Dropout(dropout) def forward(self, xyz, img_features): """ Args: xyz: [B, N, 3] 3D坐标 img_features: [B, C, H, W] 图像特征 Returns: pos_embed: [B, C, H, W] 位置编码特征 """ B, C, H, W = img_features.shape N = xyz.shape[1] # 1. 对坐标进行编码 coord_embed = self.coord_proj(xyz) # [B, N, C_pos] # 2. 对图像特征进行投影，生成空间注意力权重 feat_weights = self.feature_proj(img_features) # [B, C_pos, H, W] feat_weights = F.softmax(feat_weights.view(B, self.num_pos_feats, -1), dim=-1) feat_weights = feat_weights.view(B, self.num_pos_feats, H, W) # 3. 将坐标编码与特征权重结合 # 扩展坐标编码以匹配空间维度 coord_embed = coord_embed.unsqueeze(-1).unsqueeze(-1) # [B, N, C_pos, 1, 1] feat_weights = feat_weights.unsqueeze(1) # [B, 1, C_pos, H, W] # 加权聚合：每个空间位置的编码 = sum(coord_embed * feat_weights) weighted_coord = (coord_embed * feat_weights).sum(dim=2) # [B, N, H, W] # 4. 融合特征 # 将加权坐标与原始特征拼接 feat_flat = img_features.view(B, C, -1) # [B, C, H*W] weighted_coord_flat = weighted_coord.view(B, N, -1) # [B, N, H*W] # 为每个空间位置选择最相关的坐标编码 # 计算相似度：coord_embed @ feat_weights.T similarity = torch.einsum('bnc,bchw->bnhw', coord_embed.squeeze(-1).squeeze(-1), feat_weights.squeeze(1)) topk_indices = torch.topk(similarity, k=min(3, N), dim=1)[1] # [B, 3, H, W] # 聚合top-k坐标编码 selected_coords = torch.gather( coord_embed.squeeze(-1).squeeze(-1).unsqueeze(-1).unsqueeze(-1), dim=1, index=topk_indices.unsqueeze(2) ).sum(dim=1) # [B, C_pos, H, W] # 与原始特征拼接并融合 fused_feat = torch.cat([img_features, selected_coords], dim=1) pos_embed = self.fusion(fused_feat.view(B, -1, H*W).transpose(1,2)) # [B, H*W, C_out] pos_embed = pos_embed.transpose(1,2).view(B, -1, H, W) # [B, C_out, H, W] return self.dropout(pos_embed) # 辅助函数：生成BEV网格坐标 def generate_bev_grid(x_range, y_range, z_range, device='cuda'): """ 生成BEV空间网格坐标 Args: x_range: [min_x, max_x, step] y_range: [min_y, max_y, step] z_range: [min_z, max_z, step] device: 计算设备 Returns: grid_coords: [N, 3] 网格坐标 """ xs = torch.arange(*x_range, device=device) ys = torch.arange(*y_range, device=device) zs = torch.arange(*z_range, device=device) grid_x, grid_y, grid_z = torch.meshgrid(xs, ys, zs, indexing='ij') grid_coords = torch.stack([grid_x, grid_y, grid_z], dim=-1) return grid_coords.view(-1, 3) # 使用示例 if __name__ == "__main__": # 创建BEV网格 bev_grid = generate_bev_grid( x_range=[-51.2, 51.2, 0.4], y_range=[-51.2, 51.2, 0.4], z_range=[-5.0, 3.0, 0.4] ) # [256*256*20, 3] ≈ 1.3M points # 创建位置编码器 pos_encoder = FeatureGuidedPositionEncoding( in_channels=256, out_channels=256, num_pos_feats=128 ).to('cuda') # 模拟图像特征 img_feat = torch.randn(2, 256, 32, 80).to('cuda') # [B, C, H, W] # 生成位置编码 pos_embed = pos_encoder(bev_grid.unsqueeze(0), img_feat) print(f"Position embedding shape: {pos_embed.shape}") # [2, 256, 32, 80] 

这个实现包含了PETRV2v2的核心创新——特征引导机制。它不是简单地将坐标编码加到特征上，而是让图像特征"指导"位置编码的生成过程。具体来说，图像特征被用来生成空间注意力权重，这些权重决定了哪些坐标编码应该在哪些空间位置上被强调。

4.3 位置编码的集成应用

在模型训练流程中，位置编码需要与图像特征正确融合。以下是典型的集成方式：

# models/petrv2/backbone.py import torch import torch.nn as nn from .position_encoding import FeatureGuidedPositionEncoding class PETRV2Backbone(nn.Module): """PETRV2主干网络""" def __init__(self, img_backbone_cfg=None, img_neck_cfg=None, position_encoding_cfg=None): super().__init__() # 图像主干网络（如ResNet50） self.img_backbone = build_backbone(img_backbone_cfg) # 图像特征颈部（如FPN） self.img_neck = build_neck(img_neck_cfg) # 位置编码器 self.position_encoder = FeatureGuidedPositionEncoding(**position_encoding_cfg) # 特征投影层（将图像特征映射到合适维度） self.feature_proj = nn.Conv2d(256, 256, 1) def forward(self, img_inputs, img_metas): """ Args: img_inputs: [B, N_cam, C, H, W] 多视角图像 img_metas: 图像元信息列表 Returns: img_features: [B, C, H, W] 编码后的图像特征 """ B, N_cam, C, H, W = img_inputs.shape # 1. 提取多视角特征 img_features_list = [] for i in range(N_cam): # 单视角特征提取 feat = self.img_backbone(img_inputs[:, i]) feat = self.img_neck(feat)[-1] # 取最后一层特征 img_features_list.append(feat) # 2. 融合多视角特征（简单平均） img_features = torch.stack(img_features_list, dim=1).mean(dim=1) # [B, C, H, W] # 3. 应用位置编码 # 生成BEV网格坐标（这里简化，实际中应根据BEV范围生成） bev_grid = generate_bev_grid( x_range=[-51.2, 51.2, 0.4], y_range=[-51.2, 51.2, 0.4], z_range=[-5.0, 3.0, 0.4], device=img_features.device ) # 投影图像特征到合适维度 proj_features = self.feature_proj(img_features) # 生成位置编码 pos_embed = self.position_encoder(bev_grid.unsqueeze(0), proj_features) # 4. 特征融合：原始特征 + 位置编码 fused_features = img_features + pos_embed return fused_features 

这种集成方式确保了位置信息被自然地融入到特征学习过程中，而不是作为外部附加信息。模型在训练时会自动学习如何利用位置编码来提升3D检测性能。

5. 损失函数设计与实现

5.1 多任务损失函数架构

PETRV2是一个多任务模型，同时执行3D目标检测、BEV分割和3D车道线检测。因此，它的损失函数是多个子任务损失的加权和：

总损失 = λ_det × L_det + λ_seg × L_seg + λ_lane × L_lane 

其中，L_det是3D检测损失，L_seg是BEV分割损失，L_lane是车道线检测损失。权重λ用于平衡不同任务的梯度幅度。

对于3D检测任务，PETRV2采用DETR风格的端到端检测框架，使用匈牙利算法进行预测与真值的最优匹配。这避免了传统方法中复杂的anchor设计和NMS后处理。

5.2 Focal Loss + IoU Loss实现

3D检测的分类和回归需要不同的损失函数。分类使用Focal Loss解决正负样本不平衡问题，回归使用IoU Loss提高定位精度：

# models/petrv2/losses.py import torch import torch.nn as nn import torch.nn.functional as F from torch.nn import CrossEntropyLoss from mmdet.models.losses import SmoothL1Loss class FocalLoss(nn.Module): """Focal Loss实现""" def __init__(self, alpha=1.0, gamma=2.0, reduction='mean'): super().__init__() self.alpha = alpha self.gamma = gamma self.reduction = reduction def forward(self, inputs, targets): """ Args: inputs: [N, C] 预测logits targets: [N] 真值标签 Returns: loss: 标量损失值 """ ce_loss = F.cross_entropy(inputs, targets, reduction='none') pt = torch.exp(-ce_loss) focal_weight = (1 - pt) ** self.gamma if self.alpha >= 0: alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets) focal_weight = alpha_t * focal_weight loss = focal_weight * ce_loss if self.reduction == 'mean': return loss.mean() elif self.reduction == 'sum': return loss.sum() else: return loss class IoULoss(nn.Module): """3D IoU Loss实现""" def __init__(self, eps=1e-6, reduction='mean'): super().__init__() self.eps = eps self.reduction = reduction def forward(self, pred_boxes, target_boxes): """ Args: pred_boxes: [N, 7] 预测3D框 [x,y,z,l,w,h,theta] target_boxes: [N, 7] 真值3D框 Returns: loss: IoU损失值 """ # 计算3D IoU（简化版，实际中应使用更精确的3D IoU计算） # 这里使用2D BEV IoU作为近似 pred_xy = pred_boxes[:, :2] pred_lw = pred_boxes[:, 3:5] target_xy = target_boxes[:, :2] target_lw = target_boxes[:, 3:5] # 计算2D BEV交集 pred_min = pred_xy - pred_lw / 2 pred_max = pred_xy + pred_lw / 2 target_min = target_xy - target_lw / 2 target_max = target_xy + target_lw / 2 inter_min = torch.max(pred_min, target_min) inter_max = torch.min(pred_max, target_max) inter_area = torch.clamp(inter_max - inter_min, min=0).prod(dim=1) # 计算2D BEV并集 pred_area = pred_lw.prod(dim=1) target_area = target_lw.prod(dim=1) union_area = pred_area + target_area - inter_area # 计算IoU iou = inter_area / (union_area + self.eps) # IoU Loss = 1 - IoU loss = 1 - iou if self.reduction == 'mean': return loss.mean() elif self.reduction == 'sum': return loss.sum() else: return loss class PETRV2Loss(nn.Module): """PETRV2多任务损失函数""" def __init__(self, loss_cls=dict(type='FocalLoss', alpha=0.25, gamma=2.0), loss_bbox=dict(type='IoULoss'), loss_iou=dict(type='SmoothL1Loss'), loss_seg=dict(type='CrossEntropyLoss'), loss_lane=dict(type='FocalLoss'), loss_weights=dict( loss_cls=2.0, loss_bbox=0.25, loss_iou=0.25, loss_seg=1.0, loss_lane=1.0)): super().__init__() # 初始化各子任务损失 self.loss_cls = FocalLoss(**loss_cls) self.loss_bbox = IoULoss(**loss_bbox) self.loss_iou = SmoothL1Loss(**loss_iou) self.loss_seg = CrossEntropyLoss(**loss_seg) self.loss_lane = FocalLoss(**loss_lane) self.loss_weights = loss_weights def forward(self, pred_dict, target_dict): """ Args: pred_dict: 预测字典 - cls_scores: [B, N_q, C] 分类分数 - bbox_preds: [B, N_q, 7] 3D框预测 - seg_preds: [B, C_seg, H, W] BEV分割预测 - lane_preds: [B, N_lane, C_lane]