first commit

This commit is contained in:
2026-06-11 03:33:14 +08:00
commit 5f555bf342
599 changed files with 142347 additions and 0 deletions

View File

@@ -0,0 +1,274 @@
---
name: architecture-design
description: Use only when creating new registrable ML components that require Factory or Registry patterns.
version: 1.2.0
---
# 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.
## When to Use
Use this skill when:
- Creating a new Dataset class that needs `@register_dataset`
- Creating a new Model class that needs `@register_model`
- Creating a new module directory with `__init__.py` factory wiring
- Initializing a new ML project structure from scratch
- Adding new component types such as Augmentation, CollateFunction, or Metrics
## When Not to Use
Do not use this skill when:
- Modifying existing functions or methods
- Fixing bugs in existing code
- Adding helper functions or utilities
- Refactoring without adding new registrable components
- Making simple code changes to a single file
- Modifying configuration files
- Reading or understanding existing code
Key indicator: if the task does not require a `@register_*` decorator or a Factory pattern, skip this skill.
## 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:
1. Create file in `src/data_module/dataset/`
2. Use `@register_dataset("name")` decorator
3. Inherit from `torch.utils.data.Dataset`
4. 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:
1. Create file in `src/model_module/model/` or appropriate module subdirectory
2. Use `@register_model('ModelName')` decorator
3. `__init__` accepts **ONLY** `cfg` parameter - all hyperparameters come from config
4. `forward()` returns dict: `{"loss": loss, "labels": labels, "logits": logits}`
5. 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:
1. Create file in `src/data_module/augmentation/`
2. Implement transformation function
3. 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
1. Read the relevant module's factory/registry pattern
2. Check existing implementations for consistency
3. Follow the established directory structure
4. Use registration decorators for new components
### Adding New Features
1. Determine which module the feature belongs to
2. Check if similar functionality exists
3. Follow factory/registry pattern if creating new component types
4. Add configuration files if needed
5. 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

View File

@@ -0,0 +1,117 @@
"""
Data Augmentation Example
Demonstrates how to create a custom data augmentation function
following the architecture design pattern.
"""
import torch
from typing import Dict
from src.data_module.augmentation import register_augmentation
@register_augmentation("time_shift")
def time_shift(signal: torch.Tensor, max_shift: int = 10) -> torch.Tensor:
"""Randomly shift signal in time.
Args:
signal: Input signal tensor of shape (channels, time_steps)
max_shift: Maximum number of steps to shift
Returns:
Shifted signal tensor
"""
shift = torch.randint(-max_shift, max_shift + 1, (1,)).item()
return torch.roll(signal, shifts=shift, dims=-1)
@register_augmentation("amplitude_scale")
def amplitude_scale(
signal: torch.Tensor,
min_scale: float = 0.8,
max_scale: float = 1.2
) -> torch.Tensor:
"""Randomly scale signal amplitude.
Args:
signal: Input signal tensor of shape (channels, time_steps)
min_scale: Minimum scaling factor
max_scale: Maximum scaling factor
Returns:
Scaled signal tensor
"""
scale = torch.empty(1).uniform_(min_scale, max_scale).item()
return signal * scale
@register_augmentation("gaussian_noise")
def add_gaussian_noise(
signal: torch.Tensor,
mean: float = 0.0,
std: float = 0.1
) -> torch.Tensor:
"""Add Gaussian noise to signal.
Args:
signal: Input signal tensor of shape (channels, time_steps)
mean: Mean of Gaussian noise
std: Standard deviation of Gaussian noise
Returns:
Signal with added noise
"""
noise = torch.randn_like(signal) * std + mean
return signal + noise
# Example: Composed augmentation
@register_augmentation("composed")
def composed_augmentation(signal: torch.Tensor, cfg) -> torch.Tensor:
"""Apply multiple augmentations in sequence.
Args:
signal: Input signal tensor
cfg: Configuration object with augmentation parameters
Returns:
Augmented signal tensor
"""
# Apply each augmentation based on config
if cfg.augmentation.time_shift:
signal = time_shift(signal, cfg.augmentation.max_shift)
if cfg.augmentation.amplitude_scale:
signal = amplitude_scale(
signal,
cfg.augmentation.min_scale,
cfg.augmentation.max_scale
)
if cfg.augmentation.gaussian_noise:
signal = add_gaussian_noise(
signal,
cfg.augmentation.noise_mean,
cfg.augmentation.noise_std
)
return signal
# Usage in dataset class
class AugmentedDataset:
"""Example dataset with augmentation support."""
def __init__(self, cfg):
self.cfg = cfg
self.augmentation_fn = AugmentationFactory(cfg.augmentation.name)
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
# Load signal
signal = self.load_signal(idx)
# Apply augmentation (training mode only)
if self.training and self.augmentation_fn:
signal = self.augmentation_fn(signal, self.cfg)
return {"signal": signal, "label": self.labels[idx]}

View File

@@ -0,0 +1,131 @@
# Hydra Configuration Example
# This demonstrates the config structure for training pipeline
# Run with: python train.py --config-name=config_example
defaults:
- training: default
- dataset: brain_decoder
- model: transformer
- override hydra/launcher: submitit_local
# Project settings
project_name: brain_decoder
experiment_name: transformer_baseline
# Random seed
seed: 42
# Device settings
device: cuda
num_workers: 4
pin_memory: true
# Training configuration
training:
epochs: 100
batch_size: 32
learning_rate: 0.001
weight_decay: 0.0001
gradient_clip: 1.0
early_stopping:
patience: 10
min_delta: 0.001
# Optimizer settings
optimizer: adamw
optimizer_params:
betas: [0.9, 0.999]
eps: 1.0e-08
# Scheduler settings
scheduler: cosine
scheduler_params:
warmup_epochs: 10
min_lr: 1.0e-06
# Checkpoint settings
checkpoint:
save_every: 5
save_best: true
monitor: val_loss
mode: min
# Dataset configuration
dataset:
name: brain_decoder
task: movement_classification
target_size:
movement_classification: 5
reconstruction: [64, 64]
# Data paths
data_dir: ${dir.data_dir}/processed
train_split: train
val_split: val
test_split: test
# Data loading
num_channels: 64
sampling_rate: 1000
sequence_length: 1000
# Augmentation
augmentation:
name: composed
time_shift: true
amplitude_scale: true
gaussian_noise: true
max_shift: 10
min_scale: 0.8
max_scale: 1.2
noise_mean: 0.0
noise_std: 0.1
# Model configuration
model:
name: Transformer
hidden_dim: 256
num_heads: 8
num_layers: 6
dropout: 0.1
activation: gelu
# Architecture specific
encoder:
input_dim: 64
embedding_dim: 256
positional_encoding: true
decoder:
output_dim: ${dataset.target_size.${dataset.task}}
pooling: avg
# Logging configuration
logging:
logger: wandb
log_every: 10
log_grads: false
# TensorBoard
tensorboard: true
histogram: true
# W&B
wandb:
project: ${project_name}
entity: null
tags: ["baseline", "transformer"]
# Output directories
dir:
data_dir: ./data
output_dir: ./outputs
log_dir: ${dir.output_dir}/logs
checkpoint_dir: ${dir.output_dir}/checkpoints
figure_dir: ${dir.output_dir}/figures
table_dir: ${dir.output_dir}/tables
# Debug settings
debug: false
fast_dev_run: false

View File

@@ -0,0 +1,50 @@
"""
Example: Creating a Custom Dataset
This example shows how to add a new dataset following the project architecture.
"""
from torch.utils.data import Dataset
from typing import Dict
import torch
from src.data_module.dataset import register_dataset
@register_dataset("time_series")
class TimeSeriesDataset(Dataset):
"""
Time series dataset for sequence modeling.
Args:
sequences: List of time series sequences
seq_length: Fixed sequence length (pad or truncate if needed)
"""
def __init__(self, sequences: list, seq_length: int = 100):
self.sequences = sequences
self.seq_length = seq_length
def __len__(self) -> int:
return len(self.sequences)
def __getitem__(self, i: int) -> Dict[str, torch.Tensor]:
sequence = self.sequences[i]
# Pad or truncate to fixed length
if len(sequence) < self.seq_length:
padding = torch.zeros(self.seq_length - len(sequence))
sequence = torch.cat([sequence, padding])
else:
sequence = sequence[:self.seq_length]
return {
"input": sequence,
"label": sequence, # For autoencoder, etc.
"length": torch.tensor(min(len(self.sequences[i]), self.seq_length))
}
# Usage in training:
# from src.data_module.dataset import DatasetFactory
# dataset = DatasetFactory("time_series")(sequences=training_data, seq_length=128)
# dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

View File

@@ -0,0 +1,217 @@
"""
Example: Creating a Custom Model
This example shows how to add a new model following the project architecture.
IMPORTANT: Models use a config-driven pattern where __init__ only accepts cfg.
Key Requirements:
- Use @register_model('ModelName') decorator
- __init__ accepts ONLY cfg parameter
- All hyperparameters come from cfg (cfg.model.*, cfg.dataset.*, etc.)
- forward() returns dict: {"loss": loss, "labels": labels, "logits": logits}
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Optional
# Import the register_model decorator
# Location may vary: src.model_module.brain_decoder or src.model_module.model
from src.model_module.brain_decoder import register_model
@register_model('SimpleMLP')
class SimpleMLP(nn.Module):
"""
Simple Multi-Layer Perceptron for classification tasks.
Config structure ( Hydra YAML ):
model:
input_dim: 100
hidden_dim: 256
output_dim: 10
num_layers: 3
dropout: 0.1
dataset:
task: classification # Used to get target_size
target_size:
classification: 10
"""
def __init__(self, cfg):
super().__init__()
# Store config
self.cfg = cfg
# Get task info from config
self.task = cfg.dataset.task
# Build model - ALL parameters from cfg
self.input_dim = cfg.model.input_dim
self.hidden_dim = cfg.model.get('hidden_dim', 256)
self.output_dim = cfg.dataset.target_size[cfg.dataset.task]
self.num_layers = cfg.model.get('num_layers', 3)
self.dropout = cfg.model.get('dropout', 0.1)
# Build layers
layers = []
in_dim = self.input_dim
for i in range(self.num_layers):
layers.extend([
nn.Linear(in_dim, self.hidden_dim),
nn.ReLU(),
nn.Dropout(self.dropout)
])
in_dim = self.hidden_dim
# Output layer
layers.append(nn.Linear(self.hidden_dim, self.output_dim))
self.network = nn.Sequential(*layers)
# Loss function
self.loss_fn = nn.CrossEntropyLoss()
def forward(
self,
x: torch.Tensor,
labels: Optional[torch.Tensor] = None,
**kwargs
) -> Dict[str, Optional[torch.Tensor]]:
"""
Forward pass.
Args:
x: Input tensor of shape (batch_size, input_dim)
labels: Ground truth labels (optional, for training)
Returns:
Dictionary with:
- loss: Computed loss (None if labels not provided)
- labels: Ground truth labels
- logits: Model predictions
"""
logits = self.network(x)
loss = None
if labels is not None:
# Convert labels to long type if needed
if labels.dtype != torch.long:
labels = labels.long()
loss = self.loss_fn(logits, labels)
return {
"loss": loss,
"labels": labels,
"logits": logits
}
# ============================================
# Example with Training/Inference Modes
# ============================================
@register_model('SimpleMLPWithModes')
class SimpleMLPWithModes(nn.Module):
"""
MLP with separate training and inference logic.
Shows how to handle different modes using self.training.
"""
def __init__(self, cfg):
super().__init__()
self.cfg = cfg
self.task = cfg.dataset.task
self.input_dim = cfg.model.input_dim
self.hidden_dim = cfg.model.get('hidden_dim', 256)
self.output_dim = cfg.dataset.target_size[cfg.dataset.task]
self.fc_in = nn.Linear(self.input_dim, self.hidden_dim)
self.ln = nn.LayerNorm(self.hidden_dim)
self.fc_out = nn.Linear(self.hidden_dim, self.output_dim)
self.loss_fn = nn.CrossEntropyLoss()
# Test-time augmentation config
self.tta_times = cfg.model.get('tta_times', 1)
def forward(
self,
x: torch.Tensor,
labels: Optional[torch.Tensor] = None,
**kwargs
) -> Dict[str, Optional[torch.Tensor]]:
"""
Forward pass with training/inference modes.
"""
if self.training:
# Training mode
x = x.float()
x = self.fc_in(x)
x = self.ln(x)
x = F.relu(x)
logits = self.fc_out(x)
loss = None
if labels is not None:
if labels.dtype != torch.long:
labels = labels.long()
loss = self.loss_fn(logits, labels)
return {
"loss": loss,
"labels": labels,
"logits": logits
}
else:
# Inference mode with TTA
all_logits = []
with torch.no_grad():
x = x.float()
for _ in range(self.tta_times):
x_aug = x.clone()
# Apply TTA transformations here if needed
x_aug = self.fc_in(x_aug)
x_aug = self.ln(x_aug)
x_aug = F.relu(x_aug)
logits = self.fc_out(x_aug)
all_logits.append(logits)
# Average predictions
avg_logits = torch.mean(torch.stack(all_logits), dim=0)
loss = None
if labels is not None:
if labels.dtype != torch.long:
labels = labels.long()
loss = self.loss_fn(avg_logits, labels)
return {
"loss": loss,
"labels": labels,
"logits": avg_logits
}
# ============================================
# Config Example (Hydra YAML)
# ============================================
"""
# run/conf/model/simple_mlp.yaml
model:
name: SimpleMLP
input_dim: 100
hidden_dim: 256
output_dim: 10
num_layers: 3
dropout: 0.1
tta_times: 1
# Then in training pipeline:
# from src.model_module.brain_decoder import ModelFactory
# model = ModelFactory(cfg.model.name)(cfg)
"""

View File

@@ -0,0 +1,189 @@
#!/bin/bash
###############################################################################
# Training Pipeline Script
#
# This script demonstrates the standard training pipeline execution pattern.
# It handles environment setup, configuration, and execution with proper
# error handling and logging.
#
# Usage:
# ./run/pipeline/training/train.sh --config-name=config_example
###############################################################################
set -e # Exit on error
set -o pipefail # Exit on pipe failure
# Script configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
EXPERIMENT_NAME="baseline_experiment"
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color
###############################################################################
# Helper Functions
###############################################################################
log_info() {
echo -e "${GREEN}[INFO]${NC} $1"
}
log_warn() {
echo -e "${YELLOW}[WARN]${NC} $1"
}
log_error() {
echo -e "${RED}[ERROR]${NC} $1"
}
cleanup() {
log_info "Cleaning up..."
# Add cleanup logic here (e.g., kill background processes)
}
trap cleanup EXIT
###############################################################################
# Environment Setup
###############################################################################
setup_environment() {
log_info "Setting up environment..."
# Activate virtual environment if it exists
if [ -f "${PROJECT_ROOT}/.venv/bin/activate" ]; then
source "${PROJECT_ROOT}/.venv/bin/activate"
log_info "Activated virtual environment"
fi
# Check required commands
command -v python >/dev/null 2>&1 || { log_error "Python not found"; exit 1; }
# Set Python path
export PYTHONPATH="${PROJECT_ROOT}/src:${PYTHONPATH}"
log_info "PYTHONPATH set to: ${PYTHONPATH}"
}
###############################################################################
# Configuration
###############################################################################
parse_arguments() {
# Default values
CONFIG="default"
GPUS=0
SEED=42
# Parse arguments
while [[ $# -gt 0 ]]; do
case $1 in
--config-name|-c)
CONFIG="$2"
shift 2
;;
--gpus|-g)
GPUS="$2"
shift 2
;;
--seed|-s)
SEED="$2"
shift 2
;;
*)
log_warn "Unknown argument: $1"
shift
;;
esac
done
log_info "Configuration: ${CONFIG}"
log_info "GPUs: ${GPUS}"
log_info "Seed: ${SEED}"
}
###############################################################################
# Main Training Function
###############################################################################
run_training() {
log_info "Starting training..."
# Output directory for this run
OUTPUT_DIR="${PROJECT_ROOT}/outputs/${EXPERIMENT_NAME}/${TIMESTAMP}"
mkdir -p "${OUTPUT_DIR}"
log_info "Output directory: ${OUTPUT_DIR}"
# Training command with Hydra
python "${PROJECT_ROOT}/train.py" \
--config-name="${CONFIG}" \
seed=${SEED} \
dir.output_dir="${OUTPUT_DIR}" \
training.device=cuda \
hydra.output_dir="${OUTPUT_DIR}/hydra" \
hydra.run.dir="${OUTPUT_DIR}/hydra" || {
log_error "Training failed!"
exit 1
}
log_info "Training completed successfully!"
}
###############################################################################
# Post-Processing
###############################################################################
post_process() {
log_info "Post-processing results..."
# Copy logs to output directory
if [ -f "${OUTPUT_DIR}/hydra/*.log" ]; then
cp "${OUTPUT_DIR}/hydra/"*.log "${OUTPUT_DIR}/"
fi
# Generate summary
log_info "Run summary:"
log_info " Config: ${CONFIG}"
log_info " Seed: ${SEED}"
log_info " Output: ${OUTPUT_DIR}"
# Print path to best checkpoint
BEST_CHECKPOINT=$(find "${OUTPUT_DIR}" -name "best*.pt" | head -n 1)
if [ -n "${BEST_CHECKPOINT}" ]; then
log_info " Best checkpoint: ${BEST_CHECKPOINT}"
fi
}
###############################################################################
# Main Execution
###############################################################################
main() {
log_info "=========================================="
log_info "Training Pipeline"
log_info "=========================================="
# Setup
setup_environment
# Parse arguments
parse_arguments "$@"
# Run training
run_training
# Post-process
post_process
log_info "=========================================="
log_info "Pipeline completed successfully!"
log_info "=========================================="
}
# Run main function
main "$@"

View File

@@ -0,0 +1,117 @@
# Auto-Import Pattern
## Overview
The Auto-Import pattern automatically discovers and imports all submodules in a directory, ensuring all components are registered without manual imports.
## Structure
```python
# In module __init__.py (e.g., data_module/dataset/__init__.py)
import os
from src.utils.helpers import import_modules
models_dir = os.path.dirname(__file__)
import_modules(models_dir, "src.data_module.dataset")
```
## Helper Function
```python
# In src/utils/helpers.py
import os
import importlib
import pkgutil
from typing import List
def import_modules(models_dir: str, package_name: str) -> List[str]:
"""
Import all Python modules in a directory.
Args:
models_dir: Directory path to scan
package_name: Full package name for imports
Returns:
List of imported module names
"""
imported = []
for module_loader, name, ispkg in pkgutil.iter_modules([models_dir]):
if not name.startswith('_'):
full_name = f"{package_name}.{name}"
importlib.import_module(full_name)
imported.append(name)
return imported
```
## Benefits
- **Zero maintenance**: Adding new file = auto-registration
- **No遗漏**: Cannot forget to import new component
- **Consistent**: All components follow same discovery path
- **Scalable**: Works for any number of submodules
## Implementation Details
1. Scan directory for `.py` files
2. Skip files starting with `_` (private)
3. Import each module using full package path
4. Import triggers decorator registration
## Directory Structure Example
```
dataset/
├── __init__.py # Contains import_modules() call
├── simple_dataset.py # Auto-imported, registers "simple"
├── custom_dataset.py # Auto-imported, registers "custom"
└── _private.py # NOT imported (starts with _)
```
## Best Practices
- **Skip private files**: Files starting with `_` are not imported
- **Full package paths**: Use dot-notation for correct imports
- **Idempotent**: Safe to call multiple times
- **Error handling**: Import errors propagate for debugging
## Common Patterns
### Conditional Import
```python
def import_modules(models_dir: str, package_name: str, skip: List[str] = None):
skip = skip or []
for module_loader, name, ispkg in pkgutil.iter_modules([models_dir]):
if name not in skip and not name.startswith('_'):
importlib.import_module(f"{package_name}.{name}")
```
### Recursive Import
```python
def import_modules_recursive(models_dir: str, package_name: str):
"""Import modules and subpackages recursively."""
for importer, name, ispkg in pkgutil.walk_packages([models_dir], prefix=f"{package_name}."):
if not name.split('.')[-1].startswith('_'):
importlib.import_module(name)
```
### Dry-Run Mode
```python
def import_modules(models_dir: str, package_name: str, dry_run: bool = False):
if dry_run:
return [name for _, name, _ in pkgutil.iter_modules([models_dir])
if not name.startswith('_')]
# ... actual import logic
```
## Integration with Registry
The auto-import pattern is typically used WITH registry pattern:
1. **Import time**: `import_modules()` imports all files
2. **Decorator execution**: `@register_dataset()` runs
3. **Factory population**: `DATASET_FACTORY` dict populated
4. **Runtime**: `DatasetFactory()` looks up registered classes

View File

@@ -0,0 +1,234 @@
# Code Style Guidelines
## Type Annotations
Always use type hints for function signatures and class attributes:
```python
from typing import Dict, List, Optional, Tuple
import torch
def __getitem__(self, i: int) -> Dict[str, torch.Tensor]:
"""Get item by index."""
return self.data[i]
def compute_metrics(predictions: torch.Tensor, labels: torch.Tensor) -> Dict[str, float]:
"""Compute evaluation metrics."""
pass
class MyModel(nn.Module):
hidden_dim: int # Class attribute type hints
output_dim: int
def __init__(self, cfg):
self.hidden_dim: int = cfg.model.hidden_dim
self.output_dim: int = cfg.model.output_dim
```
## Import Order
Organize imports in three sections with blank lines between:
```python
# 1. Standard library imports
import os
from typing import Dict, List, Optional
from pathlib import Path
# 2. Third-party imports
import torch
import torch.nn as nn
from torch.utils.data import Dataset
import numpy as np
# 3. Local imports
from src.data_module.dataset import register_dataset
from src.utils.helpers import import_modules
from src.model_module.brain_decoder import register_model
```
## __init__.py Files
### Module __init__.py (with factory)
Contains factory/registry logic and auto-import:
```python
# src/data_module/dataset/__init__.py
import os
from typing import Dict, Callable, TypeVar
from src.utils.helpers import import_modules
T = TypeVar('T')
DATASET_FACTORY: Dict[str, type] = {}
def register_dataset(name: str) -> Callable[[T], T]:
"""Decorator to register dataset classes."""
def decorator(cls: T) -> T:
DATASET_FACTORY[name] = cls
return cls
return decorator
def DatasetFactory(data_name: str):
"""Create dataset instance by name."""
dataset = DATASET_FACTORY.get(data_name, None)
if dataset is None:
dataset = DATASET_FACTORY.get('simple')
return dataset
# Auto-import all submodules
models_dir = os.path.dirname(__file__)
import_modules(models_dir, "src.data_module.dataset")
```
### Subpackage __init__.py (can be empty)
```python
# src/data_module/augmentation/__init__.py
# Empty file - just marks as package
```
Or with exports:
```python
# src/data_module/__init__.py
from .dataset import DatasetFactory, register_dataset
from .augmentation import AugmentationFactory
```
## Naming Conventions
### Files
- **Modules**: `simple_dataset.py`, `custom_model.py`
- **Pipelines**: `training.sh`, `inference.sh`
- **Configs**: `config.yaml`, `brain_decoder.yaml`
- **Utilities**: `get_optimizer.py`, `helpers.py`, `compute_metrics.py`
### Classes and Functions
```python
# Classes: PascalCase
class SimpleDataset(Dataset):
pass
class MyCustomModel(nn.Module):
pass
# Functions and variables: snake_case
def compute_accuracy(predictions, labels):
pass
def get_optimizer(cfg):
pass
learning_rate = 0.001
batch_size = 32
```
### Constants
```python
# Constants: UPPER_SNAKE_CASE
DEFAULT_HIDDEN_DIM = 256
MAX_EPOCHS = 100
LEARNING_RATE = 0.001
```
## Docstrings
Use Google-style docstrings:
```python
def DatasetFactory(data_name: str) -> type:
"""Create dataset class by name.
Args:
data_name: Name of the dataset to create.
Returns:
Dataset class if found, otherwise simple dataset.
Raises:
ValueError: If no dataset is found and no default exists.
"""
pass
```
## Configuration-Driven Classes
Model classes must be config-driven:
```python
@register_model('MyModel')
class MyModel(nn.Module):
def __init__(self, cfg):
"""Initialize model from config.
Args:
cfg: Hydra config object with model attributes.
"""
super().__init__()
self.cfg = cfg
# ALL parameters from cfg
self.hidden_dim = cfg.model.hidden_dim
self.output_dim = cfg.dataset.target_size[cfg.dataset.task]
self.dropout = cfg.model.dropout
def forward(self, x, labels=None, **kwargs):
"""Forward pass.
Args:
x: Input tensor.
labels: Ground truth labels (training mode).
**kwargs: Additional arguments.
Returns:
Dict with loss, labels, and logits.
"""
# Implementation
return {"loss": loss, "labels": labels, "logits": logits}
```
## Error Handling
```python
def DatasetFactory(data_name: str) -> type:
"""Create dataset class by name."""
dataset = DATASET_FACTORY.get(data_name)
if dataset is None:
available = ', '.join(DATASET_FACTORY.keys())
raise ValueError(
f"Dataset '{data_name}' not found. "
f"Available: {available}"
)
return dataset
```
## Logging
```python
import logging
logger = logging.getLogger(__name__)
@register_dataset('custom')
class CustomDataset(Dataset):
def __init__(self, cfg):
self.cfg = cfg
logger.info(f"Initializing {self.__class__.__name__}")
logger.debug(f"Config: {cfg.dataset}")
```
## Code Review Checklist
- [ ] All functions have type hints
- [ ] Imports are correctly ordered
- [ ] Classes use PascalCase, functions use snake_case
- [ ] Docstrings follow Google style
- [ ] Model classes are config-driven
- [ ] Registration decorators are used
- [ ] Error messages are informative
- [ ] Logging is added for key operations

View File

@@ -0,0 +1,53 @@
# Factory Pattern
## Overview
The Factory pattern allows dynamic creation of instances without specifying the exact class. Each module uses a factory to decouple creation from usage.
## Structure
```python
# In module __init__.py (e.g., data_module/dataset/__init__.py)
DATASET_FACTORY: Dict[str, type] = {}
def DatasetFactory(data_name: str):
"""Create dataset instance by name."""
dataset = DATASET_FACTORY.get(data_name, None)
if dataset is None:
# Fallback to default
dataset = DATASET_FACTORY.get('simple')
return dataset
```
## Usage
```python
# Consumer code doesn't need to know concrete class
dataset = DatasetFactory(cfg.dataset.name)
```
## Benefits
- **Loose coupling**: Consumer doesn't import concrete classes
- **Extensibility**: Add new types without changing consumer code
- **Fallback handling**: Graceful degradation for unknown types
- **Centralized registry**: Single source of truth for available types
## Implementation Details
1. Define factory dict at module level
2. Factory function handles lookup and fallback
3. Return class (not instance) for deferred initialization
4. None result triggers fallback to default implementation
## Common Patterns
```python
# With config integration
def DatasetFactory(cfg):
data_name = cfg.dataset.name
dataset_cls = DATASET_FACTORY.get(data_name)
if dataset_cls is None:
raise ValueError(f"Unknown dataset: {data_name}")
return dataset_cls(cfg)
```

View File

@@ -0,0 +1,102 @@
# Registry Pattern
## Overview
The Registry pattern allows components to register themselves via decorators, enabling automatic discovery and centralized management of available types.
## Structure
```python
# In module __init__.py (e.g., data_module/dataset/__init__.py)
from typing import Dict, Callable, TypeVar
T = TypeVar('T')
DATASET_FACTORY: Dict[str, type] = {}
def register_dataset(name: str) -> Callable[[T], T]:
"""Decorator to register dataset classes."""
def decorator(cls: T) -> T:
DATASET_FACTORY[name] = cls
return cls
return decorator
```
## Usage
```python
# In implementation file (e.g., simple_dataset.py)
from data_module.dataset import register_dataset
@register_dataset("simple")
class SimpleDataset(Dataset):
def __init__(self, cfg):
# Implementation
pass
```
## Benefits
- **Automatic registration**: Components register themselves on import
- **Declarative**: Single decorator line replaces manual registration code
- **Import-time discovery**: Auto-import pattern finds all implementations
- **Type-safe**: Preserves original class type
## Implementation Details
1. Decorator returns the class unchanged (for immediate use)
2. Side effect: adds class to factory dict
3. Name parameter must be unique per module
4. Registration happens at module import time
## Advanced Patterns
### Registration with Config
```python
def register_model(name: str):
def decorator(cls):
MODEL_FACTORY[name] = cls
# Add config validation
cls._config_schema = getattr(cls, '_config_schema', {})
return cls
return decorator
```
### Conditional Registration
```python
def register_dataset(name: str, experimental: bool = False):
def decorator(cls):
if not experimental or cfg.enable_experimental:
DATASET_FACTORY[name] = cls
return cls
return decorator
```
### Multi-Registry
```python
# Multiple registries in one module
DATASET_FACTORY = {}
AUGMENTATION_FACTORY = {}
def register_dataset(name: str):
def decorator(cls):
DATASET_FACTORY[name] = cls
return cls
return decorator
def register_augmentation(name: str):
def decorator(fn):
AUGMENTATION_FACTORY[name] = fn
return fn
return decorator
```
## Best Practices
- **Unique names**: Use descriptive, unique registration names
- **Documentation**: Document required parameters in class docstring
- **Validation**: Validate config in `__init__`, not in decorator
- **Consistency**: Use same naming convention across modules

View File

@@ -0,0 +1,150 @@
# Detailed Directory Structure
This document provides a comprehensive breakdown of the ML project template directory structure.
## Root Level Files
| File | Purpose |
|------|---------|
| `README.md` | Project documentation, installation guide, usage examples |
| `TODO.md` | Task tracking with weekly focus and daily tasks |
| `.gitignore` | Git ignore patterns for Python, Jupyter, IDEs, logs, cache |
| `pyproject.toml` | Project configuration for build system and dependencies |
| `uv.lock` | Locked dependency versions for reproducibility |
## run/ - Execution Layer
### pipeline/
Main workflow scripts organized by stage:
| Directory | Purpose |
|-----------|---------|
| `training/` | Training execution scripts (training.sh, inference.sh) |
| `prepare_data/` | Data preparation and preprocessing pipelines |
| `analysis/` | Evaluation and analysis workflows |
### conf/
Hydra configuration files organized by module:
| Directory | Purpose |
|-----------|---------|
| `training/` | Training hyperparameters, model configs, optimizer settings |
| `dataset/` | Dataset configurations, data paths, preprocessing options |
| `model/` | Model architecture configurations |
| `prepare_data/` | Data preparation parameters |
| `analysis/` | Analysis and evaluation configurations |
| `dir/` | Directory path configurations |
| `analysis/` | Analysis-specific settings |
## src/ - Source Code Layer
### data_module/ - Data Processing Module
```
data_module/
├── __init__.py # Module exports
├── utils.py # Data-specific utility functions
├── dataset/ # Dataset implementations
│ ├── __init__.py # Dataset factory and registry
│ └── simple_dataset.py # Simple dataset example
├── augmentation/ # Data augmentation methods
│ ├── __init__.py
│ ├── mixup.py # Mixup augmentation
│ ├── random_shift.py # Random shifting
│ ├── channel_mask.py # Channel masking
│ ├── time_masking.py # Time masking
│ └── add_noise.py # Noise injection
├── collate_fn/ # Batch collation functions
│ ├── __init__.py
│ └── simple_collate_fn.py
├── compute_metrics/ # Metrics computation
│ ├── __init__.py
│ └── simple_compute_metrics.py
├── prepare_data/ # Data preparation logic
│ ├── __init__.py
│ ├── prepare_data.py
│ └── generate_yaml.py
└── data_func/ # Data utility functions
├── __init__.py
└── simple_data_func.py
```
### model_module/ - Model Module
```
model_module/
├── __init__.py # Module exports
└── model/ # Model implementations
└── [model files]
```
### trainer_module/ - Training Module
Contains training loop logic, validation, and checkpoint management.
### analysis_module/ - Analysis Module
Contains evaluation, visualization, and result analysis code.
### llm/ - LLM Module
LLM-related code and integrations.
### utils/ - Shared Utilities
```
utils/
├── __init__.py
├── helpers.py # Helper functions (import_modules, etc.)
├── logging.py # Logging configuration
├── get_optimizer.py # Optimizer factory
├── get_scheduler.py # Learning rate scheduler factory
├── get_callback.py # Training callbacks
├── get_activation.py # Activation functions
└── get_checkpoint_aggregation.py # Checkpoint handling
```
## data/ - Data Layer
Following the Cookiecutter Data Science standard:
| Directory | Purpose |
|-----------|---------|
| `raw/` | Original, immutable data dump |
| `processed/` | Cleaned, transformed data ready for use |
| `external/` | Data from third-party sources |
## outputs/ - Output Layer
| Directory | Purpose |
|-----------|---------|
| `logs/` | Training logs, tensorboard logs |
| `checkpoints/` | Model checkpoints for resuming training |
| `tables/` | Result tables, CSV outputs |
| `figures/` | Plots, visualizations, figures |
## Module Interaction Flow
```
run/pipeline/ -> src/trainer_module/ -> src/model_module/
src/data_module/ src/utils/
src/utils/
run/conf/ -> Hydra config loader -> All modules
```
## File Naming Conventions
- **Modules**: `simple_dataset.py`, `custom_model.py`
- **Pipelines**: `training.sh`, `inference.sh`
- **Configs**: `config.yaml`, dataset-specific names
- **Utilities**: Descriptive names (`get_optimizer.py`, `helpers.py`)
## Python Package Structure
Each module is a proper Python package:
- Has `__init__.py` with factory/registry logic
- Can be imported as `from src.module import Component`
- Subpackages are automatically discovered via `import_modules()`