NotACracker/COTR项目模型定制化开发指南

2025-07-04 17:19:51作者：尤辰城Agatha

模型组件概述

在NotACracker/COTR项目中，3D目标检测模型的架构通常可以分为6大核心组件，每个组件承担着不同的功能：

编码器(Encoder)：处理原始点云数据，包括体素化层(voxel layer)、体素编码器(voxel encoder)和中间编码器(middle encoder)，如HardVFE和PointPillarsScatter等
骨干网络(Backbone)：通常是全卷积网络(FCN)，用于提取特征图，如ResNet、SECOND等
颈部网络(Neck)：连接骨干网络和检测头的中间组件，如FPN、SECONDFPN等
检测头(Head)：执行特定任务的组件，如边界框预测和掩码预测
RoI提取器(RoI extractor)：从特征图中提取感兴趣区域特征，如H3DRoIHead和PartAggregationROIHead
损失函数(Loss)：检测头中用于计算损失的组件，如FocalLoss、L1Loss和GHMLoss等

自定义组件开发流程

开发新的编码器

以开发HardVFE体素特征编码器为例：

创建编码器类：在指定目录下创建新文件，定义编码器类并使用装饰器注册：

import torch.nn as nn
from ..builder import VOXEL_ENCODERS

@VOXEL_ENCODERS.register_module()
class HardVFE(nn.Module):
    def __init__(self, arg1, arg2):
        # 初始化参数
        pass
    
    def forward(self, x):
        # 实现前向传播逻辑
        pass

导入模块：可以通过修改__init__.py文件或配置文件的custom_imports实现
配置使用：在模型配置中指定新编码器类型和参数

开发新的骨干网络

以SECOND网络为例：

定义网络结构：创建新文件定义网络类并注册：

import torch.nn as nn
from ..builder import BACKBONES

@BACKBONES.register_module()
class SECOND(BaseModule):
    def __init__(self, arg1, arg2):
        # 网络结构定义
        pass
    
    def forward(self, x):
        # 特征提取逻辑
        pass

导入与配置：与编码器类似，可通过多种方式导入并在配置中指定

开发新的颈部网络

以SECONDFPN为例：

实现颈部网络：创建新文件定义网络结构：

from ..builder import NECKS

@NECKS.register_module()
class SECONDFPN(BaseModule):
    def __init__(self, in_channels, out_channels, upsample_strides):
        # 特征金字塔网络实现
        pass
    
    def forward(self, X):
        # 多尺度特征融合逻辑
        pass

配置使用：在模型配置中指定输入输出通道等参数

开发新的检测头

以PartA2检测头为例，这是一个两阶段检测器中的RoI Head：

实现边界框头：创建新文件定义bbox head：

from mmdet.models.builder import HEADS
from .bbox_head import BBoxHead

@HEADS.register_module()
class PartA2BboxHead(BaseModule):
    def __init__(self, num_classes, seg_in_channels, part_in_channels):
        # 初始化分类和回归分支
        pass
    
    def forward(self, seg_feats, part_feats):
        # 前向计算逻辑
        pass

实现RoI Head：继承Base3DRoIHead实现完整检测头：

@HEADS.register_module()
class PartAggregationROIHead(Base3DRoIHead):
    def __init__(self, semantic_head, num_classes, seg_roi_extractor=None):
        # 初始化各组件
        pass
    
    def _bbox_forward(self, seg_feats, part_feats, voxels_dict, rois):
        # 特征提取和预测逻辑
        pass

配置使用：在模型配置中详细指定各组件参数

开发新的损失函数

以自定义回归损失MyLoss为例：

实现损失函数：创建新文件定义损失计算：

from ..builder import LOSSES
from .utils import weighted_loss

@weighted_loss
def my_loss(pred, target):
    # 自定义损失计算逻辑
    return torch.abs(pred - target)

@LOSSES.register_module()
class MyLoss(nn.Module):
    def __init__(self, reduction='mean', loss_weight=1.0):
        # 初始化参数
        pass
    
    def forward(self, pred, target, weight=None):
        # 加权损失计算
        return self.loss_weight * my_loss(pred, target, ...)