Hi-FT/ERD项目模型自定义开发指南

模型组件概述

在Hi-FT/ERD项目中，模型架构被系统地划分为五个核心组件，这种模块化设计使得开发者可以灵活地替换或扩展各个部分：

1. 骨干网络(Backbone)：通常是全卷积网络(FCN)，用于提取图像特征，如ResNet、MobileNet等
2. 颈部网络(Neck)：连接骨干网络和头部网络的中间组件，如FPN(特征金字塔网络)、PAFPN等
3. 头部网络(Head)：执行特定任务的组件，如边界框预测、掩码预测等
4. ROI提取器(RoI Extractor)：从特征图中提取感兴趣区域(RoI)特征的组件，如RoI Align
5. 损失函数(Loss)：头部网络中用于计算损失的部分，如FocalLoss、L1Loss等

自定义骨干网络开发

1. 创建新骨干网络

以开发MobileNet为例，我们需要创建一个新的Python文件mobilenet.py：

import torch.nn as nn
from mmdet.registry import MODELS

@MODELS.register_module()
class MobileNet(nn.Module):
    def __init__(self, arg1, arg2):
        # 初始化网络结构
        super().__init__()
        # 定义网络层
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1)
        # 可根据需要添加更多层
        
    def forward(self, x):
        # 定义前向传播逻辑
        x = self.conv1(x)
        # 必须返回一个元组
        return (x,)

2. 注册新模块

有两种方式注册新模块：

方法一：直接修改__init__.py文件

from .mobilenet import MobileNet

方法二：通过配置文件动态导入（推荐）

custom_imports = dict(
    imports=['mmdet.models.backbones.mobilenet'],
    allow_failed_imports=False)

3. 在配置中使用新骨干

model = dict(
    backbone=dict(
        type='MobileNet',
        arg1=value1,  # 自定义参数
        arg2=value2),
    ...
)

自定义颈部网络开发

颈部网络通常用于特征融合和增强，如构建特征金字塔。

1. 定义颈部网络

以PAFPN为例：

@MODELS.register_module()
class PAFPN(nn.Module):
    def __init__(self, in_channels, out_channels, num_outs):
        super().__init__()
        # 初始化特征金字塔各层
        self.lateral_convs = nn.ModuleList()
        for in_channel in in_channels:
            self.lateral_convs.append(
                nn.Conv2d(in_channel, out_channels, 1))
        
    def forward(self, inputs):
        # 实现特征融合逻辑
        ...

2. 注册与使用

注册方式与骨干网络类似，使用时在配置中指定：

neck=dict(
    type='PAFPN',
    in_channels=[256, 512, 1024, 2048],  # 输入特征图通道数
    out_channels=256,  # 输出统一通道数
    num_outs=5)  # 输出特征图数量

自定义头部网络开发

头部网络是任务特定的组件，我们以Double Head R-CNN为例说明开发流程。

1. 定义新的边界框头部

@MODELS.register_module()
class DoubleConvFCBBoxHead(BBoxHead):
    def __init__(self, num_convs, num_fcs, **kwargs):
        super().__init__(**kwargs)
        # 初始化卷积分支和全连接分支
        self.conv_branch = nn.Sequential(...)
        self.fc_branch = nn.Sequential(...)
        
    def forward(self, x_cls, x_reg):
        # 实现双分支前向逻辑
        conv_feat = self.conv_branch(x_cls)
        fc_feat = self.fc_branch(x_reg)
        return cls_score, bbox_pred

2. 定义新的ROI头部

@MODELS.register_module()
class DoubleHeadRoIHead(StandardRoIHead):
    def __init__(self, reg_roi_scale_factor, **kwargs):
        super().__init__(**kwargs)
        self.reg_roi_scale_factor = reg_roi_scale_factor
        
    def _bbox_forward(self, x, rois):
        # 重写bbox前向传播
        bbox_cls_feats = self.extract_feat(x, rois)
        bbox_reg_feats = self.extract_feat(
            x, rois, scale_factor=self.reg_roi_scale_factor)
        return self.bbox_head(bbox_cls_feats, bbox_reg_feats)

3. 完整配置示例

model = dict(
    roi_head=dict(
        type='DoubleHeadRoIHead',
        bbox_head=dict(
            type='DoubleConvFCBBoxHead',
            num_convs=4,
            num_fcs=2,
            ...  # 其他参数
        )
    )
)

自定义损失函数开发

1. 实现新损失函数

@weighted_loss  # 装饰器实现加权
def my_loss(pred, target):
    # 实现损失计算逻辑
    return torch.abs(pred - target)

@MODELS.register_module()
class MyLoss(nn.Module):
    def __init__(self, reduction='mean', loss_weight=1.0):
        super().__init__()
        self.reduction = reduction
        self.loss_weight = loss_weight
        
    def forward(self, pred, target, weight=None):
        loss = my_loss(pred, target, weight)
        return self.loss_weight * loss

2. 使用新损失

在头部配置中指定：

loss_bbox=dict(type='MyLoss', loss_weight=1.0)

开发建议

1. 模块化设计：保持每个组件的独立性，便于替换和复用
2. 继承现有组件：尽可能继承现有实现，只重写必要部分
3. 参数化设计：通过配置文件灵活控制组件行为
4. 文档规范：为每个新组件添加清晰的文档说明
5. 测试验证：开发后需进行充分测试确保功能正确

通过遵循这些指南，开发者可以高效地为Hi-FT/ERD项目扩展新功能，同时保持代码的整洁性和可维护性。