在Ultralytics YOLO中添加/更换模块详细指南

一、添加新模块的步骤

1. 创建新模块（在nn/modules/目录下）

在ultralytics/nn/modules/目录中创建新模块文件（如my_module.py）:

import torch.nn as nn

class MyModule(nn.Module):
    """自定义模块说明"""
    
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):
        super().__init__()
        # 模块实现
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=kernel_size//2)
        self.act = nn.SiLU()
        
    def forward(self, x):
        return self.act(self.conv(x))

2. 在nn/modules/__init__.py中导出模块

# 在__init__.py中添加
from .my_module import MyModule

3. 在tasks.py的parse_model函数中注册新模块

找到parse_model函数中的模块识别部分，添加对新模块的支持：

def parse_model(d, ch, verbose=True):
    # ... [其他代码]
    
    if m in {
        Classify,
        Conv,
        # ... [其他模块]
        MyModule,  # 添加新模块
    }:
        c1, c2 = ch[f], args[0]
        # 处理参数...
        
    # ... [其他代码]

4. 在模型YAML配置中使用新模块

创建或修改模型配置文件（如yolov8-my.yaml）: