| name | architecture-design |
| description | 仅在创建需要 Factory 或 Registry 模式的新 registrable ML 组件时使用。 |
| version | 1.2.0 |
架构设计 - ML Project Template
该 skill 定义了基于模板结构的 machine learning 项目标准代码架构。在修改或扩展代码时,应遵循这些模式以维持一致性。
概览
该项目采用模块化、可扩展的架构,并明确分离关注点。每个模块(data、model、trainer、analysis)都通过 factory 和 registry patterns 独立组织,以获得最大的灵活性。
For non-ML architecture work, use these general constraints when this skill is already active: keep boundaries explicit, dependencies flowing in one direction, modules cohesive, public interfaces stable, and migration paths reversible. Prefer the smallest architecture change that solves the stated problem.
何时使用
在以下场景使用该 skill:
- 创建需要
@register_dataset 的新 Dataset class
- 创建需要
@register_model 的新 Model class
- 创建带有
__init__.py factory wiring 的新模块目录
- 从零初始化新的 ML 项目结构
- 新增 Augmentation、CollateFunction 或 Metrics 等组件类型
何时不使用
在以下场景不要使用该 skill:
- 修改已有函数或方法
- 修复现有代码中的 bug
- 添加 helper functions 或 utilities
- 仅重构而不新增 registrable components
- 对单文件做简单代码修改
- 修改配置文件
- 阅读或理解现有代码
关键判断标准:如果任务不需要 @register_* decorator 或 Factory pattern,就跳过这个 skill。
For general refactors, architecture review can still use the checklist below without importing the ML-specific factory requirements.
核心设计模式
Factory Pattern
每个模块都使用 factory 动态创建实例:
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。
Registry Pattern
组件通过 decorators 注册自身:
@register_dataset("simple")
class SimpleDataset(Dataset):
def __init__(self, data):
self.data = data
详细说明参见 references/registry_pattern.md。
Auto-Import Pattern
模块会自动发现并导入子模块:
models_dir = os.path.dirname(__file__)
import_modules(models_dir, "src.data_module.dataset")
详细说明参见 references/auto_import.md。
目录结构
project/
|- run/
| |- pipeline/ # 主工作流脚本
| | |- training/ # 训练流水线
| | |- prepare_data/ # 数据准备流水线
| | \- analysis/ # 分析流水线
| \- conf/ # Hydra 配置文件
| |- training/ # 训练配置
| |- dataset/ # 数据集配置
| |- model/ # 模型配置
| |- prepare_data/ # 数据准备配置
| \- analysis/ # 分析配置
|
|- src/
| |- data_module/ # 数据处理模块
| | |- dataset/ # 数据集实现
| | |- augmentation/ # 数据增强
| | |- collate_fn/ # Collate 函数
| | |- compute_metrics/ # 指标计算
| | |- prepare_data/ # 数据准备逻辑
| | |- data_func/ # 数据工具函数
| | \- utils.py # 模块内部工具
| |
| |- model_module/ # 模型实现
| | |- brain_decoder/ # Brain decoder 模型
| | \- model/ # 备选模型目录
| |
| |- trainer_module/ # 训练逻辑
| |- analysis_module/ # 分析与评估
| |- llm/ # LLM 相关代码
| \- utils/ # 共享工具
|
|- data/
| |- raw/ # 原始、不可变数据
| |- processed/ # 清洗与转换后的数据
| \- external/ # 第三方数据
|
|- outputs/
| |- logs/ # 训练与评估日志
| |- checkpoints/ # 模型检查点
| |- tables/ # 结果表格
| \- figures/ # 图表与可视化
|
|- pyproject.toml # 项目配置
|- uv.lock # 依赖锁文件
|- TODO.md # 任务跟踪
|- README.md # 项目文档
\- .gitignore # Git 忽略规则
更详细的目录说明参见 references/structure.md。
模块组织
创建新 Dataset
新增 dataset 时:
- 在
src/data_module/dataset/ 中创建文件
- 使用
@register_dataset("name") decorator
- 继承
torch.utils.data.Dataset
- 实现
__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]
创建新 Model
CRITICAL:模型使用 config-driven 模式
新增 model 时:
- 在
src/model_module/model/ 或对应子目录中创建文件
- 使用
@register_model('ModelName') decorator
__init__ 只能接收 cfg 参数,所有超参数都来自配置
forward() 返回 dict:{"loss": loss, "labels": labels, "logits": logits}
- 使用
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
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:
pass
else:
pass
return {"loss": loss, "labels": labels, "logits": logits}
添加数据增强
新增 augmentation 时:
- 在
src/data_module/augmentation/ 中创建文件
- 实现 transformation function
- 如有需要,通过 factory 注册
代码风格指南
完整风格指南参见 references/code_style.md。
关键原则:
- 函数签名始终使用类型标注
- 遵循 import 顺序:standard library -> third-party -> local
- 模块
__init__.py 包含 factory / registry 逻辑
- Model classes 必须是 config-driven
配置管理
项目使用 Hydra 管理配置:
- 配置文件放在
run/conf/,按模块组织
- 每个 stage(training、analysis)都有独立配置结构
- 所有配置均使用 YAML 文件
在本项目中工作时
修改代码前
- 阅读相关模块的 factory / registry pattern
- 检查现有实现,确保一致性
- 遵循既定目录结构
- 为新组件使用 registration decorators
添加新特性时
- 判断该功能属于哪个模块
- 检查是否已有类似功能
- 若创建新组件类型,遵循 factory / registry pattern
- 如有需要,补充配置文件
- 更新文档
代码审查清单
额外资源
参考文件
详细信息请查看:
references/structure.md - 带文件说明的详细目录结构
references/factory_pattern.md - Factory pattern 详解
references/registry_pattern.md - Registry pattern 详解
references/auto_import.md - Auto-import pattern 详解
references/code_style.md - 完整代码风格指南
示例文件
examples/ 中的可运行示例:
examples/custom_dataset.py - 自定义 dataset 实现
examples/custom_model.py - 自定义 model 实现
examples/augmentation_example.py - 数据增强示例
examples/config_example.yaml - 配置文件示例
examples/pipeline_example.sh - pipeline 脚本示例