architecture-design

架构设计 - ML 项目模板

此技能根据模板结构定义了机器学习项目的标准代码架构。修改或扩展代码时，请遵循这些模式以保持一致性。

概述

项目采用模块化、可扩展的架构，职责划分清晰。每个模块（数据、模型、训练器、分析）均独立组织，并使用工厂和注册表模式，以实现最大的灵活性。

适用场景

在以下情况下使用此技能：

• 创建需要 @register_dataset 的新 Dataset 类
• 创建需要 @register_model 的新 Model 类
• 创建包含 __init__.py 工厂配置的新模块目录
• 从头初始化新的 ML 项目结构
• 添加新的组件类型，例如 Augmentation、CollateFunction 或 Metrics

不适用场景

在以下情况下不要使用此技能：

• 修改现有函数或方法
• 修复现有代码中的错误
• 添加辅助函数或实用工具
• 不添加可注册组件的重构
• 对单个文件进行简单代码更改
• 修改配置文件
• 阅读或理解现有代码

关键判断依据：如果任务不需要 @register_* 装饰器或 Factory 模式，请跳过此技能。

核心设计模式

工厂模式

每个模块都使用工厂来动态创建实例：

# Example from data_module/dataset/__init__.py
DATASET_FACTORY: Dict = {}

def DatasetFactory(data_name: str):
    dataset = DATASET_FACTORY.get(data_name, None)
    if dataset is None:
        print(f"{data_name} dataset is not implementation, use simple dataset")
        dataset = DATASET_FACTORY.get('simple')
    return dataset

详细指南请参阅 references/factory_pattern.md。

注册表模式

组件通过装饰器自行注册：

# Example from data_module/dataset/simple_dataset.py
@register_dataset("simple")
class SimpleDataset(Dataset):
    def __init__(self, data):
        self.data = data

详细指南请参阅 references/registry_pattern.md。

自动导入模式

模块会自动发现并导入子模块：

# Example from data_module/dataset/__init__.py
models_dir = os.path.dirname(__file__)
import_modules(models_dir, "src.data_module.dataset")

详细指南请参阅 references/auto_import.md。

目录结构

project/
├── run/
│   ├── pipeline/            # Main workflow scripts
│   │   ├── training/        # Training pipelines
│   │   ├── prepare_data/    # Data preparation pipelines
│   │   └── analysis/        # Analysis pipelines
│   └── conf/                # Hydra configuration files
│       ├── training/        # Training configs
│       ├── dataset/         # Dataset configs
│       ├── model/           # Model configs
│       ├── prepare_data/    # Data prep configs
│       └── analysis/        # Analysis configs
│
├── src/
│   ├── data_module/         # Data processing module
│   │   ├── dataset/         # Dataset implementations
│   │   ├── augmentation/    # Data augmentation
│   │   ├── collate_fn/      # Collate functions
│   │   ├── compute_metrics/ # Metrics computation
│   │   ├── prepare_data/    # Data preparation logic
│   │   ├── data_func/       # Data utility functions
│   │   └── utils.py         # Module-specific utilities
│   │
│   ├── model_module/        # Model implementations
│   │   ├── brain_decoder/   # Brain decoder models
│   │   └── model/           # Alternative model location
│   │
│   ├── trainer_module/      # Training logic
│   ├── analysis_module/     # Analysis and evaluation
│   ├── llm/                 # LLM-related code
│   └── utils/               # Shared utilities
│
├── data/
│   ├── raw/                 # Original, immutable data
│   ├── processed/           # Cleaned, transformed data
│   └── external/            # Third-party data
│
├── outputs/
│   ├── logs/                # Training and evaluation logs
│   ├── checkpoints/         # Model checkpoints
│   ├── tables/              # Result tables
│   └── figures/             # Plots and visualizations
│
├── pyproject.toml           # Project configuration
├── uv.lock                  # Dependency lock file
├── TODO.md                  # Task tracking
├── README.md                # Project documentation
└── .gitignore               # Git ignore rules

如需了解包含文件说明的详细目录结构，请参阅 references/structure.md。

模块组织

创建新的数据集

添加新数据集时：

1. 在 src/data_module/dataset/ 中创建文件
2. 使用 @register_dataset("name") 装饰器
3. 继承 torch.utils.data.Dataset
4. 实现 __init__、__len__ 和 __getitem__

from torch.utils.data import Dataset
from typing import Dict
import torch
from src.data_module.dataset import register_dataset

@register_dataset("custom")
class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i: int) -> Dict[str, torch.Tensor]:
        return self.data[i]

创建新模型

重要：模型采用配置驱动模式

添加新模型时：

1. 在 src/model_module/model/ 或合适的模块子目录中创建文件
2. 使用 @register_model('ModelName') 装饰器
3. __init__ 仅接受 cfg 参数——所有超参数均来自配置
4. forward() 返回字典：{"loss": loss, "labels": labels, "logits": logits}
5. 使用 self.training 处理训练模式与推理模式

from src.model_module.brain_decoder import register_model

@register_model('MyModel')
class MyModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.task = cfg.dataset.task

        # ALL parameters from cfg
        self.hidden_dim = cfg.model.hidden_dim
        self.output_dim = cfg.dataset.target_size[cfg.dataset.task]

    def forward(self, x, labels=None, **kwargs):
        if self.training:
            # Training logic
            pass
        else:
            # Inference logic
            pass

        return {"loss": loss, "labels": labels, "logits": logits}

添加数据增强

添加数据增强时：

1. 在 src/data_module/augmentation/ 中创建文件
2. 实现转换函数
3. 如有需要，通过工厂进行注册

代码风格指南

全面的风格指南请参阅 references/code_style.md。

核心原则：

• 函数签名始终使用类型提示
• 遵循导入顺序：标准库 → 第三方库 → 本地模块
• 模块的 __init__.py 文件包含工厂/注册表逻辑
• 模型类必须采用配置驱动模式

配置管理

项目使用 Hydra 进行配置管理：

• run/conf/ 中的配置文件按模块组织
• 每个阶段（训练、分析）都有自己的配置结构
• 所有配置均使用 YAML 文件

在此项目中开展工作

修改代码前

1. 阅读相关模块的工厂/注册表模式
2. 检查现有实现以保持一致
3. 遵循既定的目录结构
4. 为新组件使用注册装饰器

添加新功能

1. 确定该功能所属的模块
2. 检查是否已有类似功能
3. 创建新的组件类型时，遵循工厂/注册表模式
4. 如有需要，添加配置文件
5. 更新文档

代码审查清单

• 适当使用工厂/注册表模式
• 遵循模块目录结构
• 添加了恰当的类型注解
• 导入项顺序正确
• 使用注册装饰器
• 根据需要添加配置文件

其他资源

参考文件

如需详细信息，请参阅：

• references/structure.md - 详细的目录结构及文件说明
• references/factory_pattern.md - 工厂模式详解
• references/registry_pattern.md - 注册表模式详解
• references/auto_import.md - 自动导入模式详解
• references/code_style.md - 全面的代码风格指南

示例文件

examples/ 中的可运行示例：

• examples/custom_dataset.py - 自定义数据集实现
• examples/custom_model.py - 自定义模型实现
• examples/augmentation_example.py - 数据增强示例
• examples/config_example.yaml - 配置文件示例
• examples/pipeline_example.sh - 流水线脚本示例