aboutsummaryrefslogtreecommitdiffstats
path: root/.agents/skills/architecture-design/examples/custom_model.py
diff options
context:
space:
mode:
authorsillylaird <sillyfanboy@gmail.com>2026-09-03 00:33:59 +0000
committersillylaird <sillyfanboy@gmail.com>2026-09-03 00:33:59 +0000
commit898b52edcb47bcb3e9d6106e74ca73e74ea01e70 (patch)
tree85c6ee5ad58b860144551184d4cf86b560c62b91 /.agents/skills/architecture-design/examples/custom_model.py
downloadwww-898b52edcb47bcb3e9d6106e74ca73e74ea01e70.tar.gz
www-898b52edcb47bcb3e9d6106e74ca73e74ea01e70.zip
import live www.sillylaird.ca webrootHEADmain
Diffstat (limited to '.agents/skills/architecture-design/examples/custom_model.py')
-rw-r--r--.agents/skills/architecture-design/examples/custom_model.py217
1 files changed, 217 insertions, 0 deletions
diff --git a/.agents/skills/architecture-design/examples/custom_model.py b/.agents/skills/architecture-design/examples/custom_model.py
new file mode 100644
index 0000000..0e00752
--- /dev/null
+++ b/.agents/skills/architecture-design/examples/custom_model.py
@@ -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)
+"""