first commit
This commit is contained in:
274
文档润色流和知识库构建流/claude-scholar/skills/architecture-design/SKILL.md
Normal file
274
文档润色流和知识库构建流/claude-scholar/skills/architecture-design/SKILL.md
Normal 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
|
||||
@@ -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]}
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
"""
|
||||
@@ -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 "$@"
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
```
|
||||
@@ -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
|
||||
@@ -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()`
|
||||
Reference in New Issue
Block a user