nn-architect

Files (2)

• SKILL.md
• architectures.md · references/

SKILL.md (first lines)

# NN Architect

Design, implement, and train neural networks with production-quality code and best practices.

---

## Workflow

```
1. Understand task and data
2. Select architecture (references/architectures.md)
3. Implement model
4. Implement training loop
5. Add training utilities (checkpointing, logging, scheduling)
```

---

## Step 1 -- Task Identification

| Task | Data type | Architecture family |
|---|---|---|
| Image classification | Images | CNN (ResNet, EfficientNet) or ViT |
| Object detection | Images + boxes | YOLO, Faster R-CNN |
| Text classification | Text sequences | Transformer encoder (BERT) |
| Text generation | Text sequences | Transformer decoder (GPT) |
| Tabular classification/regression | Structured | MLP, TabNet, FT-Transformer |
| Time series forecasting | Sequential numeric | LSTM, Temporal CNN, Transformer |
| Recommendation | User-item interactions | Embedding + MLP |
| Anomaly detection | Various | Autoencoder, VAE |

---

## Step 2 -- Select Architecture

Read `references/architectures.md` for detailed architecture selection guidance.

---

## Step 3 -- Implement Model

**Default framework**: PyTorch (unless user requests Keras/TensorFlow).

### PyTorch Model Template

```python
import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, dropout: float = 0.3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.BatchNorm1d(hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, output_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)
```

### Architecture Best Practices
- **Always include**: BatchNorm or LayerNorm, Dropout, residual connections (for deep nets).
- **Initialization**: PyTorch defaults (Kaiming) are fine for ReLU networks.
- **Activation functions**: ReLU for hidden layers (default), GELU for transformers, Sigmoid/Softmax only at output.
- **Output layer**: No activation for regression, Sigmoid for binary, LogSoftmax for multiclass (with NLLLoss) or raw logits (with CrossEntropyLoss).

---

## Step 4 -- Training Loop

### PyTorch Training Template

```python
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

def train_one_epoch(model, dataloader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    for batch in tqdm(dataloader, desc="Training"):
        inputs, targets = batch
        inputs, targets = inputs.to(device), targets.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

        total_loss += loss.item() * inputs.size(0)
        if outputs.dim() > 1:  # Classification
            correct += (outputs.argmax(dim=1) == targets).sum().item()
        total += inputs.size(0)

    avg_loss = total_loss / total
    accuracy = correct / total if correct else None
    return {"loss": avg_loss, "accuracy": accuracy}


@torch.no_grad()
def evaluate(model, dataloader, criterion, device):
    model.eval()
    total_loss = 0
    correct = 0
    total = 0

From tezansahu/ai-garage, GPL-3.0.

More from tezansahu/ai-garage