Architecture Design - ML Project Template

This skill defines the standard code architecture for machine learning projects based on the template structure. When modifying or extending code, follow these patterns to maintain consistency.

Overview

The project follows a modular, extensible architecture with clear separation of concerns. Each module (data, model, trainer, analysis) is independently organized using factory and registry patterns for maximum flexibility.

Core Design Patterns

Factory Pattern

Each module uses a factory to create instances dynamically:

python

# 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

For detailed guidance, refer to

references/factory_pattern.md

Registry Pattern

Components register themselves via decorators:

python

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

For detailed guidance, refer to

references/registry_pattern.md

Auto-Import Pattern

Modules automatically discover and import submodules:

python

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

For detailed guidance, refer to

references/auto_import.md

Directory Structure

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

For detailed directory structure with file descriptions, refer to

references/structure.md

Module Organization

Creating a New Dataset

When adding a new dataset:

Create file in
```
src/data_module/dataset/
```
Use
```
@register_dataset("name")
```
decorator
Inherit from
```
torch.utils.data.Dataset
```
Implement
```
__init__
```
,
```
__len__
```
,
```
__getitem__
```

python

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]

Creating a New Model

CRITICAL: Models use config-driven pattern

When adding a new model:

Create file in
```
src/model_module/model/
```
or appropriate module subdirectory
Use
```
@register_model('ModelName')
```
decorator
```
__init__
```
accepts ONLY
```
cfg
```
parameter - all hyperparameters come from config

forward()

returns dict:

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

Handle training vs inference modes using
```
self.training
```

python

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}

Adding Data Augmentation

When adding augmentation:

Create file in
```
src/data_module/augmentation/
```
Implement transformation function
Register with factory if needed

Code Style Guidelines

For comprehensive style guidelines, refer to

references/code_style.md

Key principles:

Always use type hints for function signatures
Follow import order: standard library → third-party → local
Module
```
__init__.py
```
files contain factory/registry logic
Model classes must be config-driven

Configuration Management

The project uses Hydra for configuration management:

Config files in
```
run/conf/
```
organize by module
Each stage (training, analysis) has its own config structure
Use YAML files for all configuration

When Working on This Project

Before Modifying Code

Read the relevant module's factory/registry pattern
Check existing implementations for consistency
Follow the established directory structure
Use registration decorators for new components

Adding New Features

Determine which module the feature belongs to
Check if similar functionality exists
Follow factory/registry pattern if creating new component types
Add configuration files if needed
Update documentation

Code Review Checklist

Uses factory/registry pattern appropriately
Follows module directory structure
Has proper type annotations
Imports are correctly ordered
Registration decorator is used
Configuration files are added if needed

Additional Resources

Reference Files

For detailed information, consult:

references/structure.md
- Detailed directory structure with file descriptions
references/factory_pattern.md
- Factory pattern in-depth explanation
references/registry_pattern.md
- Registry pattern in-depth explanation
references/auto_import.md
- Auto-import pattern in-depth explanation
references/code_style.md
- Comprehensive code style guidelines

Example Files

Working examples in

examples/

examples/custom_dataset.py
- Custom dataset implementation
examples/custom_model.py
- Custom model implementation
examples/augmentation_example.py
- Data augmentation example
examples/config_example.yaml
- Configuration file example
examples/pipeline_example.sh
- Pipeline script example

architecture-design

NPX Install

Tags

SKILL.md Content

Architecture Design - ML Project Template

Overview

Core Design Patterns

Factory Pattern

Registry Pattern

Auto-Import Pattern

Directory Structure

Module Organization

Creating a New Dataset

Creating a New Model

Adding Data Augmentation

Code Style Guidelines

Configuration Management

When Working on This Project

Before Modifying Code

Adding New Features

Code Review Checklist

Additional Resources

Reference Files

Example Files