我理解你的诉求：单阶段、体素化（voxel）、OpenPCDet/CenterPoint 基线，场景是悬挂式单轨、目标以小/少点/形状不规则为主（Bird/Kite/Plastic_sheet 明显难），希望在 mAP/mAR + FPS 上综合最优、且有大量中英文资料/论文/开源可参考，改动必须是模块级而不是“调几个阈值/学习率”。

你现在这版（CenterPoint + 你写的 PDVMeanVFE + intensity histogram）已经非常强：mAP 54.592、mAR 67.25、FPS 89，Drone/Balloon 很高，但 Kite 15.88、Bird 27.13、Plastic_sheet 45.64 是主要短板。这个分布通常意味着：** 瓶颈不在“整体能不能检测到目标”，而在“小目标特征被下采样稀释 + 类别之间头部冲突 + 评分/排序不匹配 IoU=0.7 指标”**。

我更推荐的下一步：优先做「高分辨率 + 分组双头 + IoU/质量感知打分」

在你给的约束下（单阶段 + 体素化 + 不做时序），我最推荐你下一步按模块做一套 CenterPoint-4x + Tiny/Normal 双头 + IoU-aware quality head，原因是它最直接对准你的痛点（Bird/Kite/Plastic_sheet）且资料/复现非常多：

1. 把 BEV 输出 stride 从 8 降到 4（Backbone_3D 模块改造）
   单阶段体素网络下采样过猛会损失结构/定位信息，这在 SA-SSD 等工作里就是核心问题之一。
   你的 tiny 类（Bird/Kite）在 IoU=0.7 下对定位非常敏感，“结构性提升”往往比继续玩 VFE 细节更有效。

2. 把 Dense Head 拆成至少 2 个任务头（Tiny vs Normal）
   你现在 5 类共用一套 head，梯度/尺度冲突会非常明显（tiny 类会被大类/易类“带偏”）。CBGS 的思路就是类均衡采样 + 分组多头，对小类提升非常常见。

3. 增加 IoU/Quality 分支做质量感知 re-score（Head 模块扩展）
   很多 one-stage 检测的 mAP 提升来自“分类分数能反映定位质量”（quality estimation + quality focal 类思想）。GFL/QFL 是经典体系。
   对你这种 IoU=0.7 严格 + 小目标框容易偏 的情况，收益通常很稳定，而且 FPS 影响小。

（可选 4）把 BEV 2D backbone 换成轻量 FPN/BiFPN 做多尺度融合
FPN/BiFPN 的论文与教程太多了，改起来也相对模块化。

方案A（首推）怎么按模块落到 OpenPCDet

下面完全按你要求的“模块改动”来拆解（不是“调调阈值”）。

A1) BACKBONE_3D：VoxelResBackBone8x → VoxelResBackBone4x（最关键一刀）

目标： 把稀疏 3D backbone 的总下采样倍率从 8x 改成 4x，让 HeightCompression 后的 BEV 特征更密，tiny 类更容易保留局部结构。

怎么做（工程上最省事）：

• 在 pcdet/models/backbones_3d/ 里复制 VoxelResBackBone8x 的实现，新建一个 VoxelResBackBone4x。
• 把最后一次 stride=2 的 downsample stage 改为 stride=1（或直接删掉最后一个 downsample），使总 stride 变成 4。

你必须配套改的配置（关键项）：

yaml
MODEL:
  BACKBONE_3D:
    NAME: VoxelResBackBone4x

  DENSE_HEAD:
    TARGET_ASSIGNER_CONFIG:
      FEATURE_MAP_STRIDE: 4

为什么这不是“调参”： 这是改变网络的特征图分辨率与 receptive field 的结构性变化，对 tiny 类经常是数量级影响。

直觉：你 voxel_size=0.1m，stride 8 时 BEV cell 对应 0.8m；stride 4 时变 0.4m。对 Bird/Kite 这种 box 尺寸小、IoU=0.7 严格的任务，“更密的语义栅格”通常直接提升 recall 与定位精度。

A2) DENSE_HEAD：单头 → Tiny/Normal 双头（用 CenterHead 的多任务能力 + 轻代码扩展）

你现在配置里：

yaml
CLASS_NAMES_EACH_HEAD: [
  ['Drone', 'Plastic_sheet', 'Kite', 'Balloon', 'Bird']
]

建议改成：

yaml
CLASS_NAMES_EACH_HEAD: [
  ['Bird', 'Kite'],                    # Tiny head
  ['Drone', 'Plastic_sheet', 'Balloon']# Normal head
]

为什么：

• tiny 类（Bird/Kite）更依赖高分辨率、不同正样本半径策略、不同 loss 权重；
• plastic_sheet 形状不规则/长条/面状，往往需要更强的定位约束（比如额外 IoU loss），不应该和 tiny 类共享同一套头。

建议你加的“小代码模块改动”（非常值）：
让每个 head 支持独立的 TARGET_ASSIGNER_CONFIG（例如不同 MIN_RADIUS/GAUSSIAN_OVERLAP/NUM_MAX_OBJS），而不是全局一套。这个改动位置一般在 pcdet/models/dense_heads/center_head.py 的 target 生成逻辑里：把原来单一 self.model_cfg.TARGET_ASSIGNER_CONFIG 改成 list（TARGET_ASSIGNER_CONFIG_EACH_HEAD）并按 head index 取值。

这一步非常符合 CBGS 论文的精神：分组 + 多头来缓解类别不均衡与形状差异。

A3) DENSE_HEAD：增加 IoU/Quality 分支（模块级 head 扩展）+ 质量感知重打分

你当前 SEPARATE_HEAD_CFG 是：

yaml
HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']

建议扩成：

yaml
HEAD_ORDER: ['center', 'center_z', 'dim', 'rot', 'iou']
HEAD_DICT:
  ...
  'iou': {'out_channels': 1, 'num_conv': 2}

训练怎么监督 iou：

• 对每个正样本位置，按 CenterHead 的 decode 逻辑解码出预测 3D box；
• 和它匹配的 GT box 计算 BEV IoU 或 3D IoU（只算正样本，开销可控）；
• 用这个 IoU 当作质量标签监督 iou 分支（L1/BCE 都可）。

推理怎么用：

• 用 score = heatmap_score * (iou_pred)^α（或 sqrt(heatmap_score * iou_pred)）替代单纯 heatmap_score 排序，再做 NMS。
• 这种“分类分数包含定位质量”的思想是 GFL/QFL 体系的核心之一。

这一步对你的 Bird/Kite/Plastic_sheet 很关键：
它们经常会出现“分数看起来不错，但 box 偏一点就 IoU 掉到 0.7 以下”的预测，IoU-aware re-score 能把这类预测自动降权，mAP 更稳。

A4) 回归 loss 增强：加一个 IoU 类回归 loss（模块级 loss 扩展）

你现在回归主要是 L1（center/dim/rot）。建议在正样本上加一个 BEV IoU loss（或借鉴 GIoU/DIoU 思路）让优化目标更贴近评估指标 IoU。

重点：不需要对所有候选都算 IoU，只对正样本算，成本可控。

方案B（更快更稳、几乎不掉 FPS）：CBGS + 双头 + class-wise NMS（pipeline/后处理模块）

如果你想先做一个风险最低的版本（基本不动 backbone）：

1. CBGS（数据采样模块）：显著缓解 Kite/Bird 类样本少导致的欠拟合。
2. 双头分组：同上（Tiny/Normal）。
3. class-wise NMS（后处理模块）：避免不同类别互相抑制；你现在 MULTI_CLASSES_NMS: False 有潜在风险（尤其 Bird/Kite 容易被其它类误检压掉）。
4. 需要更“软”的抑制时，用 Soft-NMS（模块更替） 替代 hard NMS 也很常见。

这套对 FPS 影响很小，但对小类 mAP 往往能有明显改善。

方案C（更激进、精度潜力更大但代码量更高）：训练期加点级辅助监督

如果你愿意加训练用的辅助分支（推理可关闭），可借鉴：

• PV-RCNN 的 point-voxel 抽象/点级监督：能帮助 backbone 学会从稀疏点里“抓目标”。
• SA-SSD 的结构感知思想：单阶段里用结构信息提升定位精度。

但这条路线工程量更大，我建议你在方案A跑通后再上。

最终结论：你下一步我建议先改哪一个？

我建议你优先落地 方案A 的前三项（最划算）：

1. VoxelResBackBone4x（stride 4）
2. Tiny/Normal 双 head（Bird+Kite / Drone+Plastic_sheet+Balloon）
3. IoU/Quality head + re-score（GFL/QFL 风格）

这三项组合是“结构性解决 tiny 类”的主路线：

• stride 4 给 tiny 类更多空间信息（单阶段下采样问题在 SA-SSD 等里被反复强调）。
• 双头 + 类均衡是 CBGS 的核心思想之一。
• 质量感知评分能显著改善 IoU=0.7 下的排序与 mAP。

你后续“从哪里找怎么改”的资料/开源依据（中英文都很多）

这些方向之所以推荐，是因为论文、复现、博客、代码都非常多：

• CenterPoint（baseline）：
• CBGS（采样 + 分组多头）：
• FPN / BiFPN（多尺度融合）：
• GFL/QFL（质量分支 + 质量感知分类）：
• PV-RCNN（点级辅助监督/点-体素融合经典）：
• SA-SSD（单阶段结构信息与定位提升）：
• （备选框架思路）VoxelNeXt（完全稀疏、速度-精度 tradeoff）：
• （可选增强）SE/CBAM/DCNv2：

如果你希望我把“方案A”进一步落到你这份 cfg里（不是泛泛而谈），我可以直接给你：

• 一份完整的 cfg 草案（包含 VoxelResBackBone4x + 双 head + iou head + per-head assigner 字段结构）
• 以及 OpenPCDet 里需要新建/修改的文件清单（到类名/函数名级别），让你按文件逐个实现，不会迷路。

下面我按方案A（CenterPoint + 体素化 + 单阶段）给你“一步一步”的代码改动 + 配置文件（可直接在 OpenPCDet 上落地）。
目标是：(1) XY stride 从 8 → 4 提升小目标分辨率，(2) Bird/Kite 与其它类分头训练减少任务冲突，(3) 加 IoU/质量分支做 re-score，提升 IoU=0.7 下的 mAP。

重要说明：你贴的 center_head.py 版本其实已经原生支持 iou 分支、USE_IOU_TO_RECTIFY_SCORE、IOU_REG_LOSS（在 get_loss 里就有相关逻辑），所以IoU 分支这步基本只需要“加 head + 开开关”即可，不需要你自己重写 IoU loss。你只需要按下面把配置补齐即可。你贴的 gitee 版本 center_head.py 在 get_loss 中确实有 calculate_iou_loss_centerhead 和 USE_IOU_TO_RECTIFY_SCORE 的逻辑。

Step 1：新增 3D Backbone：VoxelResBackBone4x（核心结构改动）

1.1 为什么这样改

你现在 voxel_size = [0.1,0.1,0.2]，X=160m、Y=40m。

• stride=8 → BEV 网格大约是 (160/0.1/8=200) × (40/0.1/8=50)
• stride=4 → 变成 400 × 100
  这对 Bird/Kite 这种“小目标 + 少点 + IoU=0.7 严格”的定位非常关键。

但要注意一个坑：stride 改了以后，Z 方向如果下采样不够，HeightCompression 输出通道会变大（C×D），会导致 2D backbone 输入通道不匹配。
所以我们要做到：

• XY 总 stride=4
• Z 下采样保持跟原 8x 类似，确保 HeightCompression 输出还是 256 通道（C=128, D≈2）

最稳的做法：把 conv4 的 stride 改为 (2,1,1)（只在 Z 下采样，不再在 XY 下采样）。

1.2 代码：在 spconv_backbone.py 增加一个新类

文件：pcdet/models/backbones_3d/spconv_backbone.py

做法：复制现有的 VoxelResBackBone8x 类，重命名为 VoxelResBackBone4x，然后按下面完整代码替换/新增（你也可以直接对照你当前 8x 类，仅改动标注的三处）。

下面这个类写法遵循 OpenPCDet 的常见风格：使用已有的 post_act_block、SparseBasicBlock、spconv_utils.replace_feature 等（这些在同文件里本来就有）。

python
# ========== 在 pcdet/models/backbones_3d/spconv_backbone.py 中新增 ==========
from functools import partial
import torch
import torch.nn as nn

from ...utils.spconv_utils import replace_feature, spconv

# 这里假设你文件里已经有 post_act_block、SparseBasicBlock
# 如果你原文件的 VoxelResBackBone8x 用的就是这些，那直接可用


class VoxelResBackBone4x(nn.Module):
    """
    XY stride: 4x (用于小目标)
    Z stride: 保持与 8x backbone 类似，保证 HeightCompression 后通道仍≈256
    """
    def __init__(self, model_cfg, input_channels, grid_size, voxel_size, point_cloud_range, **kwargs):
        super().__init__()
        self.model_cfg = model_cfg
        self.sparse_shape = grid_size[::-1] + [1, 0, 0]
        self.voxel_size = voxel_size
        self.point_cloud_range = point_cloud_range

        norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
        block = post_act_block

        self.conv_input = spconv.SparseSequential(
            spconv.SubMConv3d(input_channels, 16, 3, padding=1, bias=False, indice_key='subm1'),
            norm_fn(16),
            nn.ReLU(),
        )

        self.conv1 = spconv.SparseSequential(
            SparseBasicBlock(16, 16, norm_fn=norm_fn, indice_key='res1'),
            SparseBasicBlock(16, 16, norm_fn=norm_fn, indice_key='res1'),
        )

        # stride=2: XY, Z 都下采样一次
        self.conv2 = spconv.SparseSequential(
            block(16, 32, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
            SparseBasicBlock(32, 32, norm_fn=norm_fn, indice_key='res2'),
            SparseBasicBlock(32, 32, norm_fn=norm_fn, indice_key='res2'),
        )

        # stride=2: XY, Z 再下采样一次 (此时 XY stride=4)
        self.conv3 = spconv.SparseSequential(
            block(32, 64, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3', conv_type='spconv'),
            SparseBasicBlock(64, 64, norm_fn=norm_fn, indice_key='res3'),
            SparseBasicBlock(64, 64, norm_fn=norm_fn, indice_key='res3'),
        )

        # ========== 关键改动1：conv4 只在 Z 下采样，不在 XY 下采样 ==========
        # 原 8x：stride=2 (Z,Y,X 都下采样) -> XY stride=8
        # 现在 4x：stride=(2,1,1) -> XY stride 保持 4，Z 继续下采样
        # padding 建议保持 (0,1,1)（很多 OpenPCDet 配置就是这样，避免 Z padding）
        self.conv4 = spconv.SparseSequential(
            block(64, 128, 3, norm_fn=norm_fn, stride=(2, 1, 1), padding=(0, 1, 1),
                  indice_key='spconv4', conv_type='spconv'),
            SparseBasicBlock(128, 128, norm_fn=norm_fn, indice_key='res4'),
            SparseBasicBlock(128, 128, norm_fn=norm_fn, indice_key='res4'),
        )

        last_pad = self.model_cfg.get('last_pad', 0)

        # 与 8x 一样：只在 Z 再下采样一次，把 D 压到 ~2，保证 C*D≈256
        self.conv_out = spconv.SparseSequential(
            spconv.SparseConv3d(128, 128, (3, 1, 1), stride=(2, 1, 1), padding=last_pad,
                                bias=False, indice_key='spconv_down2'),
            norm_fn(128),
            nn.ReLU(),
        )

        self.num_point_features = 128

    def forward(self, batch_dict):
        voxel_features, voxel_coords = batch_dict['voxel_features'], batch_dict['voxel_coords']
        batch_size = batch_dict['batch_size']

        input_sp_tensor = spconv.SparseConvTensor(
            features=voxel_features,
            indices=voxel_coords.int(),
            spatial_shape=self.sparse_shape,
            batch_size=batch_size
        )

        x = self.conv_input(input_sp_tensor)
        x_conv1 = self.conv1(x)
        x_conv2 = self.conv2(x_conv1)
        x_conv3 = self.conv3(x_conv2)
        x_conv4 = self.conv4(x_conv3)

        out = self.conv_out(x_conv4)

        # ========== 关键改动2：encoded stride 改成 4 ==========
        batch_dict['encoded_spconv_tensor'] = out
        batch_dict['encoded_spconv_tensor_stride'] = 4

        # multi-scale 给其他模块用（如果你未来加 FPN/aux loss 也方便）
        batch_dict['multi_scale_3d_features'] = {
            'x_conv1': x_conv1,
            'x_conv2': x_conv2,
            'x_conv3': x_conv3,
            'x_conv4': x_conv4,
        }
        # ========== 关键改动3：x_conv4 stride 改成 4 ==========
        batch_dict['multi_scale_3d_strides'] = {
            'x_conv1': 1,
            'x_conv2': 2,
            'x_conv3': 4,
            'x_conv4': 4,
        }

        return batch_dict

Step 2：注册 backbone（让 cfg 里 NAME 能找到）

文件：pcdet/models/backbones_3d/__init__.py

把 VoxelResBackBone4x import 进来，并加到 __all__ 字典里。

python
# 原来可能是：
# from .spconv_backbone import VoxelBackBone8x, VoxelResBackBone8x

from .spconv_backbone import VoxelBackBone8x, VoxelResBackBone8x, VoxelResBackBone4x

__all__ = {
    'VoxelBackBone8x': VoxelBackBone8x,
    'VoxelResBackBone8x': VoxelResBackBone8x,
    'VoxelResBackBone4x': VoxelResBackBone4x,   # 新增这一行
    # ... 其他不动
}

Step 3：让 CenterHead 支持“每个 head 一套 target assigner”（模块级改动）

你要做 Tiny/Normal 双 head，并且希望 tiny head 用不同的 MIN_RADIUS/GAUSSIAN_OVERLAP，这就需要 CenterHead 在 assign_targets 里按 head 取不同 config。

文件：pcdet/models/dense_heads/center_head.py

修改函数：assign_targets

你现在的原始逻辑是所有 head 都用：

python
target_assigner_cfg = self.model_cfg.TARGET_ASSIGNER_CONFIG
...
feature_map_stride=target_assigner_cfg.FEATURE_MAP_STRIDE
gaussian_overlap=target_assigner_cfg.GAUSSIAN_OVERLAP
min_radius=target_assigner_cfg.MIN_RADIUS

我们改成：如果 cfg 里提供 TARGET_ASSIGNER_CONFIG_EACH_HEAD（list），就按 head_idx 取；否则回退到默认。

把下面 patch 贴进去（你只需要改 assign_targets 这一段即可）：

python
# ====== 在 assign_targets 里：target_assigner_cfg = ... 后面加 ======
target_assigner_cfg = self.model_cfg.TARGET_ASSIGNER_CONFIG
target_assigner_cfg_each_head = self.model_cfg.get('TARGET_ASSIGNER_CONFIG_EACH_HEAD', None)

def _cfg_get(cfg, key, default=None):
    if cfg is None:
        return default
    if isinstance(cfg, dict):
        return cfg.get(key, default)
    if hasattr(cfg, 'get'):
        return cfg.get(key, default)
    return getattr(cfg, key, default)

# ====== 把原来的 for idx, cur_class_names ... 改成 ======
for head_idx, cur_class_names in enumerate(self.class_names_each_head):

    # 每个 head 自己的 assigner cfg（没有就用默认）
    cur_assigner_cfg = target_assigner_cfg_each_head[head_idx] if target_assigner_cfg_each_head is not None else target_assigner_cfg

    feature_map_stride = _cfg_get(cur_assigner_cfg, 'FEATURE_MAP_STRIDE', target_assigner_cfg.FEATURE_MAP_STRIDE)
    num_max_objs = _cfg_get(cur_assigner_cfg, 'NUM_MAX_OBJS', target_assigner_cfg.NUM_MAX_OBJS)
    gaussian_overlap = _cfg_get(cur_assigner_cfg, 'GAUSSIAN_OVERLAP', target_assigner_cfg.GAUSSIAN_OVERLAP)
    min_radius = _cfg_get(cur_assigner_cfg, 'MIN_RADIUS', target_assigner_cfg.MIN_RADIUS)

    heatmap_list, target_boxes_list, inds_list, masks_list, target_boxes_src_list = [], [], [], [], []

    for bs_idx in range(batch_size):
        cur_gt_boxes = gt_boxes[bs_idx]
        gt_class_names = all_names[cur_gt_boxes[:, -1].cpu().long().numpy()]

        gt_boxes_single_head = []
        # 注意：避免 idx 变量覆盖，这里用 gt_idx
        for gt_idx, name in enumerate(gt_class_names):
            if name not in cur_class_names:
                continue

            temp_box = cur_gt_boxes[gt_idx]
            temp_box[-1] = cur_class_names.index(name) + 1
            gt_boxes_single_head.append(temp_box[None, :])

        if len(gt_boxes_single_head) == 0:
            gt_boxes_single_head = cur_gt_boxes[:0, :]
        else:
            gt_boxes_single_head = torch.cat(gt_boxes_single_head, dim=0)

        heatmap, ret_boxes, inds, mask, ret_boxes_src = self.assign_target_of_single_head(
            num_classes=len(cur_class_names),
            gt_boxes=gt_boxes_single_head.cpu(),
            feature_map_size=feature_map_size,
            feature_map_stride=feature_map_stride,
            num_max_objs=num_max_objs,
            gaussian_overlap=gaussian_overlap,
            min_radius=min_radius,
        )

        heatmap_list.append(heatmap.to(gt_boxes_single_head.device))
        target_boxes_list.append(ret_boxes.to(gt_boxes_single_head.device))
        inds_list.append(inds.to(gt_boxes_single_head.device))
        masks_list.append(mask.to(gt_boxes_single_head.device))
        target_boxes_src_list.append(ret_boxes_src.to(gt_boxes_single_head.device))

    ret_dict['heatmaps'].append(torch.stack(heatmap_list, dim=0))
    ret_dict['target_boxes'].append(torch.stack(target_boxes_list, dim=0))
    ret_dict['inds'].append(torch.stack(inds_list, dim=0))
    ret_dict['masks'].append(torch.stack(masks_list, dim=0))
    ret_dict['target_boxes_src'].append(torch.stack(target_boxes_src_list, dim=0))

你贴出来的 center_head.py 里 assign_targets 的原代码结构就是上面这种（for head 循环里又把 idx 重用了一次）。
这个 patch 做的是“模块级能力增强”，后面 cfg 才能真正生效。

Step 4：给你一份“方案A最终版 cfg”（可直接跑）

下面是基于你现有配置，做的最小侵入修改版本（保留你 PDVMeanVFE、2D backbone 等），仅加入方案A的三块：

• VoxelResBackBone4x
• 双 head（Tiny/Normal）
• iou head + rectify score
• per-head target assigner（需要 Step3 的代码支持）

你保存为：tools/cfgs/custom_models/centerpoint_pdv_4x_2head_iou.yaml（路径随你）

yaml
CLASS_NAMES: ['Drone', 'Plastic_sheet', 'Kite', 'Balloon', 'Bird']

DATA_CONFIG:
    _BASE_CONFIG_: /media/jd/4997BB1603CFE2C4/lw/Openpcdet/tools/cfgs/dataset_configs/custom_dataset.yaml

MODEL:
    NAME: Test

    VFE:
        NAME: PDVMeanVFE
        USE_XY_RANGE: True
        DENSITY_LOG1P: True
        RANGE_SCALE: 160.0
        INTENSITY_IDX: 3
        NUM_HIST_BINS: 8
        HIST_MIN: 0.0
        HIST_MAX: 1.0

    BACKBONE_3D:
        NAME: VoxelResBackBone4x   # <<< 改这里（Step1）

    MAP_TO_BEV:
        NAME: HeightCompression
        NUM_BEV_FEATURES: 256      # 仍保持 256（我们 backbone 设计就是为了不破这个）

    BACKBONE_2D:
        NAME: BaseBEVBackbone
        LAYER_NUMS: [5, 5]
        LAYER_STRIDES: [1, 2]
        NUM_FILTERS: [128, 256]
        UPSAMPLE_STRIDES: [1, 2]
        NUM_UPSAMPLE_FILTERS: [256, 256]

    DENSE_HEAD:
        NAME: CenterHead
        CLASS_AGNOSTIC: False

        # <<< 改这里：双 head（Tiny / Normal）
        CLASS_NAMES_EACH_HEAD: [
            ['Bird', 'Kite'],
            ['Drone', 'Plastic_sheet', 'Balloon']
        ]

        SHARED_CONV_CHANNEL: 64
        USE_BIAS_BEFORE_NORM: True
        NUM_HM_CONV: 2

        SEPARATE_HEAD_CFG:
            # 注意：不要把 iou 放进 HEAD_ORDER（否则回归维度和 code_weights 会错）
            HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
            HEAD_DICT: {
                'center': {'out_channels': 2, 'num_conv': 2},
                'center_z': {'out_channels': 1, 'num_conv': 2},
                'dim': {'out_channels': 3, 'num_conv': 2},
                'rot': {'out_channels': 2, 'num_conv': 2},

                # <<< 新增：IoU / quality 分支（Step4）
                'iou': {'out_channels': 1, 'num_conv': 2}
            }

        # <<< 改这里：stride 8 -> 4
        TARGET_ASSIGNER_CONFIG:
            FEATURE_MAP_STRIDE: 4
            NUM_MAX_OBJS: 500
            GAUSSIAN_OVERLAP: 0.1
            MIN_RADIUS: 2
            DENSE_REG: 1

        # <<< 新增：每个 head 自己的 assigner（需要 Step3 的代码支持）
        TARGET_ASSIGNER_CONFIG_EACH_HEAD:
            - FEATURE_MAP_STRIDE: 4
              NUM_MAX_OBJS: 500
              GAUSSIAN_OVERLAP: 0.35   # tiny：更“尖”的监督（更小 radius）
              MIN_RADIUS: 1
              DENSE_REG: 1
            - FEATURE_MAP_STRIDE: 4
              NUM_MAX_OBJS: 500
              GAUSSIAN_OVERLAP: 0.10   # normal：保持你原本风格
              MIN_RADIUS: 2
              DENSE_REG: 1

        # <<< 可选但建议：训练期加 IoU reg（不影响推理FPS）
        IOU_REG_LOSS: True

        LOSS_CONFIG:
            LOSS_WEIGHTS: {
                'cls_weight': 1.0,
                'loc_weight': 1.0,
                'code_weights': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
            }

        POST_PROCESSING:
            SCORE_THRESH: 0.1
            POST_CENTER_LIMIT_RANGE: [-160, -75.2, -5.0, 160, 75.2, 8.0]
            MAX_OBJ_PER_SAMPLE: 500

            # <<< 新增：用 iou 预测纠正 score（质量感知排序）
            USE_IOU_TO_RECTIFY_SCORE: True

            # 对应 CLASS_NAMES 顺序：Drone, Plastic_sheet, Kite, Balloon, Bird
            # alpha 越大，越依赖 iou（对 tiny 类一般设大点更稳）
            IOU_RECTIFIER: [0.50, 0.60, 0.70, 0.30, 0.70]

            NMS_CONFIG:
                MULTI_CLASSES_NMS: False
                NMS_TYPE: nms_gpu
                NMS_THRESH: 0.01
                NMS_PRE_MAXSIZE: 4096
                NMS_POST_MAXSIZE: 500

    POST_PROCESSING:
        RECALL_THRESH_LIST: [0.3, 0.5, 0.7]
        OUTPUT_RAW_SCORE: False
        EVAL_METRIC: kitti

OPTIMIZATION:
    BATCH_SIZE_PER_GPU: 4
    NUM_EPOCHS: 80

    OPTIMIZER: adam_onecycle
    LR: 0.01
    WEIGHT_DECAY: 0.01
    MOMENTUM: 0.9

    MOMS: [0.95, 0.85]
    PCT_START: 0.4
    DIV_FACTOR: 10
    DECAY_STEP_LIST: [35, 45]
    LR_DECAY: 0.1
    LR_CLIP: 0.0000001

    LR_WARMUP: False
    WARMUP_EPOCH: 1
    GRAD_NORM_CLIP: 35

Step 5：你改完后必须做的 3 个 sanity check（不然很容易“以为跑了，实际错 stride”）

1. 确认编码 stride=4
   在训练/推理 log 或者临时 print 一下（在 HeightCompression.forward 之后）：

• batch_dict['spatial_features_stride'] 应该是 4
• batch_dict['spatial_features'].shape 的 C 应该仍然是 256（否则你的 Z 下采样没对齐，会爆通道）

2. 确认 CenterHead 的 FEATURE_MAP_STRIDE=4 生效
   否则你会出现“可视化位置全偏、mAP 反而下降”。

3. 确认 iou 分支真的在输出里
   看 pred_dict 是否包含 iou 键；如果没有，说明 SEPARATE_HEAD_CFG.HEAD_DICT 没加对或你的 OpenPCDet 版本没这块逻辑（但你贴的版本是有的）。

你下一步如果想继续提升 Bird/Kite（在方案A基础上）

你先把上面 A 跑通，拿到新 mAP（特别看 Bird/Kite 是否上升）。
如果 Bird/Kite 还是明显低，我建议你下一步只加一个模块（不动其它）：

• 把 NMS 换成 class_specific_nms（后处理模块级改动）：小目标很容易被其它类压掉；这一步经常对小类 recall 有帮助。
  但这属于方案A之外，我先不强行加进来，避免一次改太多。

如果你愿意，我也可以再给你一份**“对照实验清单（A0/A1/A2/A3）”**：每一步单独开关，让你快速定位到底是 stride=4、双 head、还是 iou re-score 哪个带来的提升最大。

收到。你这份 center_head.py 跟我前面给的 patch 版本不一致，所以你会觉得“不对”。下面我按你贴出来的原始代码，给出一个可直接替换的修改版（只改必要地方）：

• ✅ 支持 TARGET_ASSIGNER_CONFIG_EACH_HEAD（每个 head 一套 assigner：stride / overlap / min_radius / max_objs）
• ✅ 修复你原代码里 idx 变量被重复覆盖 的坑（外层 head idx、内层 gt idx）
• ✅ 顺手修复一个隐藏 bug：当 IOU_REG_LOSS=True 但 pred_dict 没有 iou 时，batch_box_preds_for_iou 可能未定义（虽然你通常会加 iou head，但这里让它更鲁棒）

你只需要把下面整个文件内容替换掉原 pcdet/models/dense_heads/center_head.py（或至少替换相同位置的函数实现）。

✅ 修改后的 center_head.py（与你贴的版本对齐）

python
import copy
import numpy as np
import torch
import torch.nn as nn
from torch.nn.init import kaiming_normal_
from ..model_utils import model_nms_utils
from ..model_utils import centernet_utils
from ...utils import loss_utils
from functools import partial


class SeparateHead(nn.Module):
    def __init__(self, input_channels, sep_head_dict, init_bias=-2.19, use_bias=False, norm_func=None):
        super().__init__()
        self.sep_head_dict = sep_head_dict

        for cur_name in self.sep_head_dict:
            output_channels = self.sep_head_dict[cur_name]['out_channels']
            num_conv = self.sep_head_dict[cur_name]['num_conv']

            fc_list = []
            for k in range(num_conv - 1):
                fc_list.append(nn.Sequential(
                    nn.Conv2d(input_channels, input_channels, kernel_size=3, stride=1, padding=1, bias=use_bias),
                    nn.BatchNorm2d(input_channels) if norm_func is None else norm_func(input_channels),
                    nn.ReLU()
                ))
            fc_list.append(nn.Conv2d(input_channels, output_channels, kernel_size=3, stride=1, padding=1, bias=True))
            fc = nn.Sequential(*fc_list)
            if 'hm' in cur_name:
                fc[-1].bias.data.fill_(init_bias)
            else:
                for m in fc.modules():
                    if isinstance(m, nn.Conv2d):
                        kaiming_normal_(m.weight.data)
                        if hasattr(m, "bias") and m.bias is not None:
                            nn.init.constant_(m.bias, 0)

            self.__setattr__(cur_name, fc)

    def forward(self, x):
        ret_dict = {}
        for cur_name in self.sep_head_dict:
            ret_dict[cur_name] = self.__getattr__(cur_name)(x)
        return ret_dict


class CenterHead(nn.Module):
    def __init__(self, model_cfg, input_channels, num_class, class_names, grid_size, point_cloud_range, voxel_size,
                 predict_boxes_when_training=True):
        super().__init__()
        self.model_cfg = model_cfg
        self.num_class = num_class
        self.grid_size = grid_size
        self.point_cloud_range = point_cloud_range
        self.voxel_size = voxel_size
        self.feature_map_stride = self.model_cfg.TARGET_ASSIGNER_CONFIG.get('FEATURE_MAP_STRIDE', None)

        self.class_names = class_names
        self.class_names_each_head = []
        self.class_id_mapping_each_head = []

        for cur_class_names in self.model_cfg.CLASS_NAMES_EACH_HEAD:
            self.class_names_each_head.append([x for x in cur_class_names if x in class_names])
            cur_class_id_mapping = torch.from_numpy(np.array(
                [self.class_names.index(x) for x in cur_class_names if x in class_names]
            )).cuda()
            self.class_id_mapping_each_head.append(cur_class_id_mapping)

        total_classes = sum([len(x) for x in self.class_names_each_head])
        assert total_classes == len(self.class_names), f'class_names_each_head={self.class_names_each_head}'

        norm_func = partial(nn.BatchNorm2d, eps=self.model_cfg.get('BN_EPS', 1e-5), momentum=self.model_cfg.get('BN_MOM', 0.1))
        self.shared_conv = nn.Sequential(
            nn.Conv2d(
                input_channels, self.model_cfg.SHARED_CONV_CHANNEL, 3, stride=1, padding=1,
                bias=self.model_cfg.get('USE_BIAS_BEFORE_NORM', False)
            ),
            norm_func(self.model_cfg.SHARED_CONV_CHANNEL),
            nn.ReLU(),
        )

        self.heads_list = nn.ModuleList()
        self.separate_head_cfg = self.model_cfg.SEPARATE_HEAD_CFG
        for head_idx, cur_class_names in enumerate(self.class_names_each_head):
            cur_head_dict = copy.deepcopy(self.separate_head_cfg.HEAD_DICT)
            cur_head_dict['hm'] = dict(out_channels=len(cur_class_names), num_conv=self.model_cfg.NUM_HM_CONV)
            self.heads_list.append(
                SeparateHead(
                    input_channels=self.model_cfg.SHARED_CONV_CHANNEL,
                    sep_head_dict=cur_head_dict,
                    init_bias=-2.19,
                    use_bias=self.model_cfg.get('USE_BIAS_BEFORE_NORM', False),
                    norm_func=norm_func
                )
            )
        self.predict_boxes_when_training = predict_boxes_when_training
        self.forward_ret_dict = {}
        self.build_losses()

    def build_losses(self):
        self.add_module('hm_loss_func', loss_utils.FocalLossCenterNet())
        self.add_module('reg_loss_func', loss_utils.RegLossCenterNet())

    def assign_target_of_single_head(
            self, num_classes, gt_boxes, feature_map_size, feature_map_stride, num_max_objs=500,
            gaussian_overlap=0.1, min_radius=2
    ):
        heatmap = gt_boxes.new_zeros(num_classes, feature_map_size[1], feature_map_size[0])
        ret_boxes = gt_boxes.new_zeros((num_max_objs, gt_boxes.shape[-1] - 1 + 1))
        inds = gt_boxes.new_zeros(num_max_objs).long()
        mask = gt_boxes.new_zeros(num_max_objs).long()
        ret_boxes_src = gt_boxes.new_zeros(num_max_objs, gt_boxes.shape[-1])
        ret_boxes_src[:gt_boxes.shape[0]] = gt_boxes

        x, y, z = gt_boxes[:, 0], gt_boxes[:, 1], gt_boxes[:, 2]
        coord_x = (x - self.point_cloud_range[0]) / self.voxel_size[0] / feature_map_stride
        coord_y = (y - self.point_cloud_range[1]) / self.voxel_size[1] / feature_map_stride
        coord_x = torch.clamp(coord_x, min=0, max=feature_map_size[0] - 0.5)
        coord_y = torch.clamp(coord_y, min=0, max=feature_map_size[1] - 0.5)
        center = torch.cat((coord_x[:, None], coord_y[:, None]), dim=-1)
        center_int = center.int()
        center_int_float = center_int.float()

        dx, dy, dz = gt_boxes[:, 3], gt_boxes[:, 4], gt_boxes[:, 5]
        dx = dx / self.voxel_size[0] / feature_map_stride
        dy = dy / self.voxel_size[1] / feature_map_stride

        radius = centernet_utils.gaussian_radius(dx, dy, min_overlap=gaussian_overlap)
        radius = torch.clamp_min(radius.int(), min=min_radius)

        for k in range(min(num_max_objs, gt_boxes.shape[0])):
            if dx[k] <= 0 or dy[k] <= 0:
                continue

            if not (0 <= center_int[k][0] <= feature_map_size[0] and 0 <= center_int[k][1] <= feature_map_size[1]):
                continue

            cur_class_id = (gt_boxes[k, -1] - 1).long()
            centernet_utils.draw_gaussian_to_heatmap(heatmap[cur_class_id], center[k], radius[k].item())

            inds[k] = center_int[k, 1] * feature_map_size[0] + center_int[k, 0]
            mask[k] = 1

            ret_boxes[k, 0:2] = center[k] - center_int_float[k].float()
            ret_boxes[k, 2] = z[k]
            ret_boxes[k, 3:6] = gt_boxes[k, 3:6].log()
            ret_boxes[k, 6] = torch.cos(gt_boxes[k, 6])
            ret_boxes[k, 7] = torch.sin(gt_boxes[k, 6])
            if gt_boxes.shape[1] > 8:
                ret_boxes[k, 8:] = gt_boxes[k, 7:-1]

        return heatmap, ret_boxes, inds, mask, ret_boxes_src

    # ===================== 关键修改：assign_targets 支持 per-head assigner =====================
    def assign_targets(self, gt_boxes, feature_map_size=None, **kwargs):
        feature_map_size = feature_map_size[::-1]  # [H, W] ==> [x, y]
        target_assigner_cfg = self.model_cfg.TARGET_ASSIGNER_CONFIG

        # 新增：每个 head 独立配置（list[dict/EasyDict]）
        target_assigner_cfg_each_head = self.model_cfg.get('TARGET_ASSIGNER_CONFIG_EACH_HEAD', None)
        if target_assigner_cfg_each_head is not None:
            assert len(target_assigner_cfg_each_head) == len(self.class_names_each_head), \
                f"TARGET_ASSIGNER_CONFIG_EACH_HEAD 长度({len(target_assigner_cfg_each_head)}) " \
                f"必须等于 head 数({len(self.class_names_each_head)})"

        # 新增：如果外部传了 feature_map_stride（你 forward 里传了 spatial_features_2d_strides），优先用它
        runtime_stride = kwargs.get('feature_map_stride', None)
        if runtime_stride is not None:
            # 有些地方可能传 list/tuple，这里只取 int
            if isinstance(runtime_stride, (list, tuple)):
                runtime_stride = int(runtime_stride[0])
            else:
                runtime_stride = int(runtime_stride)

        def _get_cfg_value(cfg, key, default):
            # cfg 可能是 dict / EasyDict
            if cfg is None:
                return default
            if isinstance(cfg, dict):
                return cfg.get(key, default)
            # EasyDict / edict 都支持 getattr
            return getattr(cfg, key, default)

        batch_size = gt_boxes.shape[0]
        ret_dict = {
            'heatmaps': [],
            'target_boxes': [],
            'inds': [],
            'masks': [],
            'heatmap_masks': [],
            'target_boxes_src': [],
        }

        all_names = np.array(['bg', *self.class_names])

        for head_idx, cur_class_names in enumerate(self.class_names_each_head):
            # per-head cfg
            cur_cfg = target_assigner_cfg_each_head[head_idx] if target_assigner_cfg_each_head is not None else target_assigner_cfg

            cur_feature_map_stride = _get_cfg_value(cur_cfg, 'FEATURE_MAP_STRIDE', target_assigner_cfg.FEATURE_MAP_STRIDE)
            cur_num_max_objs = _get_cfg_value(cur_cfg, 'NUM_MAX_OBJS', target_assigner_cfg.NUM_MAX_OBJS)
            cur_gaussian_overlap = _get_cfg_value(cur_cfg, 'GAUSSIAN_OVERLAP', target_assigner_cfg.GAUSSIAN_OVERLAP)
            cur_min_radius = _get_cfg_value(cur_cfg, 'MIN_RADIUS', target_assigner_cfg.MIN_RADIUS)

            # runtime stride 优先（避免 stride=4 backbone 但 cfg 写错导致全偏）
            if runtime_stride is not None:
                cur_feature_map_stride = runtime_stride

            heatmap_list, target_boxes_list, inds_list, masks_list, target_boxes_src_list = [], [], [], [], []
            for bs_idx in range(batch_size):
                cur_gt_boxes = gt_boxes[bs_idx]
                gt_class_names = all_names[cur_gt_boxes[:, -1].cpu().long().numpy()]

                gt_boxes_single_head = []
                # 修复：不要复用 idx（原代码会覆盖外层 idx）
                for gt_idx, name in enumerate(gt_class_names):
                    if name not in cur_class_names:
                        continue
                    temp_box = cur_gt_boxes[gt_idx].clone()
                    temp_box[-1] = cur_class_names.index(name) + 1
                    gt_boxes_single_head.append(temp_box[None, :])

                if len(gt_boxes_single_head) == 0:
                    gt_boxes_single_head = cur_gt_boxes[:0, :]
                else:
                    gt_boxes_single_head = torch.cat(gt_boxes_single_head, dim=0)

                heatmap, ret_boxes, inds, mask, ret_boxes_src = self.assign_target_of_single_head(
                    num_classes=len(cur_class_names),
                    gt_boxes=gt_boxes_single_head.cpu(),
                    feature_map_size=feature_map_size,
                    feature_map_stride=cur_feature_map_stride,
                    num_max_objs=cur_num_max_objs,
                    gaussian_overlap=cur_gaussian_overlap,
                    min_radius=cur_min_radius,
                )

                heatmap_list.append(heatmap.to(gt_boxes_single_head.device))
                target_boxes_list.append(ret_boxes.to(gt_boxes_single_head.device))
                inds_list.append(inds.to(gt_boxes_single_head.device))
                masks_list.append(mask.to(gt_boxes_single_head.device))
                target_boxes_src_list.append(ret_boxes_src.to(gt_boxes_single_head.device))

            ret_dict['heatmaps'].append(torch.stack(heatmap_list, dim=0))
            ret_dict['target_boxes'].append(torch.stack(target_boxes_list, dim=0))
            ret_dict['inds'].append(torch.stack(inds_list, dim=0))
            ret_dict['masks'].append(torch.stack(masks_list, dim=0))
            ret_dict['target_boxes_src'].append(torch.stack(target_boxes_src_list, dim=0))

        return ret_dict

    def sigmoid(self, x):
        y = torch.clamp(x.sigmoid(), min=1e-4, max=1 - 1e-4)
        return y

    # ===================== 小修复：iou_reg_loss 时 batch_box_preds_for_iou 必然存在 =====================
    def get_loss(self):
        pred_dicts = self.forward_ret_dict['pred_dicts']
        target_dicts = self.forward_ret_dict['target_dicts']

        tb_dict = {}
        loss = 0

        for head_idx, pred_dict in enumerate(pred_dicts):
            pred_dict['hm'] = self.sigmoid(pred_dict['hm'])
            hm_loss = self.hm_loss_func(pred_dict['hm'], target_dicts['heatmaps'][head_idx])
            hm_loss *= self.model_cfg.LOSS_CONFIG.LOSS_WEIGHTS['cls_weight']

            target_boxes = target_dicts['target_boxes'][head_idx]
            pred_boxes = torch.cat([pred_dict[head_name] for head_name in self.separate_head_cfg.HEAD_ORDER], dim=1)

            reg_loss = self.reg_loss_func(
                pred_boxes, target_dicts['masks'][head_idx], target_dicts['inds'][head_idx], target_boxes
            )
            loc_loss = (reg_loss * reg_loss.new_tensor(self.model_cfg.LOSS_CONFIG.LOSS_WEIGHTS['code_weights'])).sum()
            loc_loss = loc_loss * self.model_cfg.LOSS_CONFIG.LOSS_WEIGHTS['loc_weight']

            loss += hm_loss + loc_loss
            tb_dict['hm_loss_head_%d' % head_idx] = hm_loss.item()
            tb_dict['loc_loss_head_%d' % head_idx] = loc_loss.item()

            if ('iou' in pred_dict) or self.model_cfg.get('IOU_REG_LOSS', False):
                batch_box_preds = centernet_utils.decode_bbox_from_pred_dicts(
                    pred_dict=pred_dict,
                    point_cloud_range=self.point_cloud_range,
                    voxel_size=self.voxel_size,
                    feature_map_stride=self.feature_map_stride
                )  # (B, H, W, 7 or 9)

                # 无论有没有 iou head，都定义 batch_box_preds_for_iou，避免未定义
                batch_box_preds_for_iou = batch_box_preds.permute(0, 3, 1, 2).contiguous()

                if 'iou' in pred_dict:
                    iou_loss = loss_utils.calculate_iou_loss_centerhead(
                        iou_preds=pred_dict['iou'],
                        batch_box_preds=batch_box_preds_for_iou.clone().detach(),
                        mask=target_dicts['masks'][head_idx],
                        ind=target_dicts['inds'][head_idx],
                        gt_boxes=target_dicts['target_boxes_src'][head_idx]
                    )
                    loss += iou_loss
                    tb_dict['iou_loss_head_%d' % head_idx] = iou_loss.item()

                if self.model_cfg.get('IOU_REG_LOSS', False):
                    iou_reg_loss = loss_utils.calculate_iou_reg_loss_centerhead(
                        batch_box_preds=batch_box_preds_for_iou,
                        mask=target_dicts['masks'][head_idx],
                        ind=target_dicts['inds'][head_idx],
                        gt_boxes=target_dicts['target_boxes_src'][head_idx]
                    )
                    if target_dicts['masks'][head_idx].sum().item() != 0:
                        iou_reg_loss = iou_reg_loss * self.model_cfg.LOSS_CONFIG.LOSS_WEIGHTS['loc_weight']
                        loss += iou_reg_loss
                        tb_dict['iou_reg_loss_head_%d' % head_idx] = iou_reg_loss.item()
                    else:
                        loss += (batch_box_preds_for_iou * 0.).sum()
                        tb_dict['iou_reg_loss_head_%d' % head_idx] = (batch_box_preds_for_iou * 0.).sum()

        tb_dict['rpn_loss'] = loss.item()
        return loss, tb_dict

    def generate_predicted_boxes(self, batch_size, pred_dicts):
        post_process_cfg = self.model_cfg.POST_PROCESSING
        post_center_limit_range = torch.tensor(post_process_cfg.POST_CENTER_LIMIT_RANGE).cuda().float()

        ret_dict = [{
            'pred_boxes': [],
            'pred_scores': [],
            'pred_labels': [],
        } for k in range(batch_size)]
        for head_idx, pred_dict in enumerate(pred_dicts):
            batch_hm = pred_dict['hm'].sigmoid()
            batch_center = pred_dict['center']
            batch_center_z = pred_dict['center_z']
            batch_dim = pred_dict['dim'].exp()
            batch_rot_cos = pred_dict['rot'][:, 0].unsqueeze(dim=1)
            batch_rot_sin = pred_dict['rot'][:, 1].unsqueeze(dim=1)
            batch_vel = pred_dict['vel'] if 'vel' in self.separate_head_cfg.HEAD_ORDER else None
            batch_iou = (pred_dict['iou'] + 1) * 0.5 if 'iou' in pred_dict else None

            final_pred_dicts = centernet_utils.decode_bbox_from_heatmap(
                heatmap=batch_hm, rot_cos=batch_rot_cos, rot_sin=batch_rot_sin,
                center=batch_center, center_z=batch_center_z, dim=batch_dim, vel=batch_vel, iou=batch_iou,
                point_cloud_range=self.point_cloud_range, voxel_size=self.voxel_size,
                feature_map_stride=self.feature_map_stride,
                K=post_process_cfg.MAX_OBJ_PER_SAMPLE,
                circle_nms=(post_process_cfg.NMS_CONFIG.NMS_TYPE == 'circle_nms'),
                score_thresh=post_process_cfg.SCORE_THRESH,
                post_center_limit_range=post_center_limit_range
            )

            for k, final_dict in enumerate(final_pred_dicts):
                final_dict['pred_labels'] = self.class_id_mapping_each_head[head_idx][final_dict['pred_labels'].long()]

                if post_process_cfg.get('USE_IOU_TO_RECTIFY_SCORE', False) and 'pred_iou' in final_dict:
                    pred_iou = torch.clamp(final_dict['pred_iou'], min=0, max=1.0)
                    IOU_RECTIFIER = final_dict['pred_scores'].new_tensor(post_process_cfg.IOU_RECTIFIER)
                    final_dict['pred_scores'] = torch.pow(
                        final_dict['pred_scores'], 1 - IOU_RECTIFIER[final_dict['pred_labels']]
                    ) * torch.pow(
                        pred_iou, IOU_RECTIFIER[final_dict['pred_labels']]
                    )

                if post_process_cfg.NMS_CONFIG.NMS_TYPE not in ['circle_nms', 'class_specific_nms']:
                    selected, selected_scores = model_nms_utils.class_agnostic_nms(
                        box_scores=final_dict['pred_scores'],
                        box_preds=final_dict['pred_boxes'],
                        nms_config=post_process_cfg.NMS_CONFIG,
                        score_thresh=None
                    )

                elif post_process_cfg.NMS_CONFIG.NMS_TYPE == 'class_specific_nms':
                    selected, selected_scores = model_nms_utils.class_specific_nms(
                        box_scores=final_dict['pred_scores'],
                        box_preds=final_dict['pred_boxes'],
                        box_labels=final_dict['pred_labels'],
                        nms_config=post_process_cfg.NMS_CONFIG,
                        score_thresh=post_process_cfg.NMS_CONFIG.get('SCORE_THRESH', None)
                    )
                elif post_process_cfg.NMS_CONFIG.NMS_TYPE == 'circle_nms':
                    raise NotImplementedError

                final_dict['pred_boxes'] = final_dict['pred_boxes'][selected]
                final_dict['pred_scores'] = selected_scores
                final_dict['pred_labels'] = final_dict['pred_labels'][selected]

                ret_dict[k]['pred_boxes'].append(final_dict['pred_boxes'])
                ret_dict[k]['pred_scores'].append(final_dict['pred_scores'])
                ret_dict[k]['pred_labels'].append(final_dict['pred_labels'])

        for k in range(batch_size):
            ret_dict[k]['pred_boxes'] = torch.cat(ret_dict[k]['pred_boxes'], dim=0)
            ret_dict[k]['pred_scores'] = torch.cat(ret_dict[k]['pred_scores'], dim=0)
            ret_dict[k]['pred_labels'] = torch.cat(ret_dict[k]['pred_labels'], dim=0) + 1

        return ret_dict

    @staticmethod
    def reorder_rois_for_refining(batch_size, pred_dicts):
        num_max_rois = max([len(cur_dict['pred_boxes']) for cur_dict in pred_dicts])
        num_max_rois = max(1, num_max_rois)
        pred_boxes = pred_dicts[0]['pred_boxes']

        rois = pred_boxes.new_zeros((batch_size, num_max_rois, pred_boxes.shape[-1]))
        roi_scores = pred_boxes.new_zeros((batch_size, num_max_rois))
        roi_labels = pred_boxes.new_zeros((batch_size, num_max_rois)).long()

        for bs_idx in range(batch_size):
            num_boxes = len(pred_dicts[bs_idx]['pred_boxes'])
            rois[bs_idx, :num_boxes, :] = pred_dicts[bs_idx]['pred_boxes']
            roi_scores[bs_idx, :num_boxes] = pred_dicts[bs_idx]['pred_scores']
            roi_labels[bs_idx, :num_boxes] = pred_dicts[bs_idx]['pred_labels']
        return rois, roi_scores, roi_labels

    def forward(self, data_dict):
        spatial_features_2d = data_dict['spatial_features_2d']
        x = self.shared_conv(spatial_features_2d)

        pred_dicts = []
        for head in self.heads_list:
            pred_dicts.append(head(x))

        if self.training:
            target_dict = self.assign_targets(
                data_dict['gt_boxes'],
                feature_map_size=spatial_features_2d.size()[2:],
                feature_map_stride=data_dict.get('spatial_features_2d_strides', None)
            )
            self.forward_ret_dict['target_dicts'] = target_dict

        self.forward_ret_dict['pred_dicts'] = pred_dicts

        if not self.training or self.predict_boxes_when_training:
            pred_dicts = self.generate_predicted_boxes(
                data_dict['batch_size'], pred_dicts
            )

            if self.predict_boxes_when_training:
                rois, roi_scores, roi_labels = self.reorder_rois_for_refining(data_dict['batch_size'], pred_dicts)
                data_dict['rois'] = rois
                data_dict['roi_scores'] = roi_scores
                data_dict['roi_labels'] = roi_labels
                data_dict['has_class_labels'] = True
            else:
                data_dict['final_box_dicts'] = pred_dicts

        return data_dict

✅ 配套 cfg 里你要加的字段（对应这个修改）

在你的 DENSE_HEAD 里加入：

yaml
CLASS_NAMES_EACH_HEAD: [
  ['Bird', 'Kite'],
  ['Drone', 'Plastic_sheet', 'Balloon']
]

TARGET_ASSIGNER_CONFIG:
  FEATURE_MAP_STRIDE: 4
  NUM_MAX_OBJS: 500
  GAUSSIAN_OVERLAP: 0.1
  MIN_RADIUS: 2
  DENSE_REG: 1

TARGET_ASSIGNER_CONFIG_EACH_HEAD:
  - FEATURE_MAP_STRIDE: 4
    NUM_MAX_OBJS: 500
    GAUSSIAN_OVERLAP: 0.35
    MIN_RADIUS: 1
  - FEATURE_MAP_STRIDE: 4
    NUM_MAX_OBJS: 500
    GAUSSIAN_OVERLAP: 0.10
    MIN_RADIUS: 2

并且你的 SEPARATE_HEAD_CFG.HEAD_DICT 加上 iou（你这份代码原生支持）：

yaml
SEPARATE_HEAD_CFG:
  HEAD_ORDER: ['center', 'center_z', 'dim', 'rot']
  HEAD_DICT: {
    'center': {'out_channels': 2, 'num_conv': 2},
    'center_z': {'out_channels': 1, 'num_conv': 2},
    'dim': {'out_channels': 3, 'num_conv': 2},
    'rot': {'out_channels': 2, 'num_conv': 2},
    'iou': {'out_channels': 1, 'num_conv': 2}
  }

IOU_REG_LOSS: True

POST_PROCESSING:
  USE_IOU_TO_RECTIFY_SCORE: True
  IOU_RECTIFIER: [0.50, 0.60, 0.70, 0.30, 0.70]