Advanced Training Example

Complete training pipeline with all features.

---

Full-Featured Example

"""
Advanced VAE Training with All Features

Demonstrates:
- Experiment tracking
- Callbacks
- Learning rate scheduling
- Early stopping
- Gradient accumulation
- Mixed precision
- Evaluation metrics
"""

import torch
import torch.optim as optim
from frameworm import Trainer, Config, get_model
from frameworm.experiment import Experiment
from frameworm.training.callbacks import (
    EarlyStopping,
    LearningRateScheduler,
    ModelCheckpoint,
    TensorBoardLogger
)
from frameworm.metrics import MetricEvaluator

# Load configuration
config = Config('config.yaml')

# Get data (from previous example)
train_loader, val_loader = get_data(config)

# Create model
model = get_model('vae')(config)

# Optimizer with weight decay
optimizer = optim.AdamW(
    model.parameters(),
    lr=config.training.lr,
    weight_decay=0.01
)

# Learning rate scheduler
scheduler = optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=config.training.epochs
)

# Create experiment
experiment = Experiment(
    name='vae-advanced',
    config=config,
    tags=['vae', 'mnist', 'advanced'],
    description='Advanced VAE training with all features'
)

# Create trainer
trainer = Trainer(
    model=model,
    optimizer=optimizer,
    device='cuda'
)

# Enable multi-GPU
trainer.enable_data_parallel(device_ids=[0, 1, 2, 3])

# Enable mixed precision (2-3x speedup)
trainer.enable_mixed_precision()

# Enable gradient accumulation (simulate larger batch)
trainer.enable_gradient_accumulation(accumulation_steps=4)

# Add callbacks
trainer.add_callback(
    EarlyStopping(
        monitor='val_loss',
        patience=10,
        mode='min'
    )
)

trainer.add_callback(
    LearningRateScheduler(scheduler)
)

trainer.add_callback(
    ModelCheckpoint(
        filepath='checkpoints/best.pt',
        monitor='val_loss',
        save_best_only=True
    )
)

trainer.add_callback(
    TensorBoardLogger(log_dir='logs')
)

# Create metric evaluator
evaluator = MetricEvaluator(
    metrics=['fid', 'is'],
    real_data=val_loader,
    device='cuda'
)

# Set evaluator (auto-evaluate every 5 epochs)
trainer.set_evaluator(evaluator, eval_every=5)

# Train with experiment tracking
with experiment:
    trainer.set_experiment(experiment)
    trainer.train(
        train_loader=train_loader,
        val_loader=val_loader,
        epochs=config.training.epochs
    )

# Final evaluation
final_metrics = evaluator.evaluate(model, num_samples=50000)

print(f"Training complete!")
print(f"Experiment ID: {experiment.experiment_id}")
print(f"Final FID: {final_metrics['fid']:.2f}")
print(f"Final IS: {final_metrics['is']:.2f}")

---

Performance Comparison

Configuration	Time/Epoch	GPU Memory	Final FID
Basic	45s	8GB	25.3
+ Multi-GPU (4x)	15s	6GB/GPU	25.1
+ Mixed Precision	8s	3GB/GPU	25.2
+ Grad Accumulation	10s	2GB/GPU	24.8

---

Tips

1. Start Simple - Add features incrementally
2. Monitor Memory - Use trainer.print_gpu_memory()
3. Track Everything - Experiments make debugging easier
4. Use Callbacks - Modular and reusable
5. Evaluate Regularly - Catch issues early EOF

cat > docs/examples/production-deployment.md << 'EOF'

Production Deployment Example

End-to-end deployment pipeline.

---

Complete Deployment Workflow

Step 1: Train Production Model

# train_production.py

from frameworm import Trainer, Config, get_model
from frameworm.experiment import Experiment

config = Config('config_production.yaml')
model = get_model('vae')(config)

with Experiment(name='vae-production-v1', config=config) as exp:
    trainer = Trainer(model, optimizer, device='cuda')
    trainer.set_experiment(exp)
    trainer.train(train_loader, val_loader, epochs=200)

print(f"Production model trained: {exp.experiment_id}")

Step 2: Export Model

# Export to multiple formats
frameworm export \
  experiments/vae-production-v1/checkpoints/best.pt \
  --format all \
  --quantize \
  --benchmark

Output: ✓ TorchScript saved: exported/model.pt (145.2 MB) ✓ ONNX saved: exported/model.onnx (144.8 MB) ✓ Quantized model saved: exported/model_quant.pt (37.1 MB) Inference Benchmark: TorchScript: 12.3 ± 0.5 ms (81 inferences/sec) ONNX Runtime: 8.7 ± 0.3 ms (115 inferences/sec) Quantized: 6.2 ± 0.2 ms (161 inferences/sec)

Step 3: Create API Server

# server.py

from frameworm.deployment import ModelServer

server = ModelServer(
    model_path='exported/model_quant.pt',  # Use quantized for speed
    device='cuda'
)

server.run(host='0.0.0.0', port=8000, workers=4)

Step 4: Containerize

# Dockerfile

FROM python:3.10-slim

# Install dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy model and code
COPY exported/model_quant.pt /app/model.pt
COPY server.py /app/

WORKDIR /app

# Expose port
EXPOSE 8000

# Health check
HEALTHCHECK --interval=30s --timeout=10s \
    CMD curl -f http://localhost:8000/health || exit 1

# Run server
CMD ["python", "server.py"]

Build and run:

docker build -t vae-server:v1 .
docker run -p 8000:8000 vae-server:v1

Step 5: Kubernetes Deployment

# k8s/deployment.yaml

apiVersion: apps/v1
kind: Deployment
metadata:
  name: vae-server
spec:
  replicas: 3
  selector:
    matchLabels:
      app: vae-server
  template:
    metadata:
      labels:
        app: vae-server
    spec:
      containers:
      - name: server
        image: vae-server:v1
        ports:
        - containerPort: 8000
        resources:
          requests:
            memory: "2Gi"
            cpu: "1"
            nvidia.com/gpu: "1"
          limits:
            memory: "4Gi"
            cpu: "2"
            nvidia.com/gpu: "1"
        livenessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 30
        readinessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 10

---
apiVersion: v1
kind: Service
metadata:
  name: vae-service
spec:
  selector:
    app: vae-server
  ports:
  - port: 80
    targetPort: 8000
  type: LoadBalancer

---
apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: vae-hpa
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: vae-server
  minReplicas: 2
  maxReplicas: 10
  metrics:
  - type: Resource
    resource:
      name: cpu
      target:
        type: Utilization
        averageUtilization: 70

Deploy:

kubectl apply -f k8s/deployment.yaml
kubectl get services  # Get external IP

Step 6: Monitoring

# prometheus-config.yaml

scrape_configs:
  - job_name: 'vae-server'
    static_configs:
      - targets: ['vae-service:8000']
    metrics_path: '/metrics'

Step 7: Load Testing

# load_test.py

import requests
import concurrent.futures
import time

def make_request():
    response = requests.post(
        'http://your-service-url/predict',
        json={'data': [[...]]}
    )
    return response.elapsed.total_seconds()

# Test with 100 concurrent requests
with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor:
    futures = [executor.submit(make_request) for _ in range(1000)]
    times = [f.result() for f in futures]

print(f"Average latency: {sum(times)/len(times)*1000:.2f}ms")
print(f"95th percentile: {sorted(times)[int(len(times)*0.95)]*1000:.2f}ms")
print(f"Throughput: {len(times)/sum(times):.0f} req/s")

---

Production Checklist

• [x] Model trained on full dataset
• [x] Exported to optimized format (ONNX/quantized)
• [x] API server with health checks
• [x] Containerized with Docker
• [x] Deployed to Kubernetes
• [x] Auto-scaling configured
• [x] Monitoring setup
• [x] Load tested

---

Cost Optimization

Setup	Cost/Month	Throughput
Single GPU (V100)	$300	1000 req/s
3x CPU (quantized)	$150	1200 req/s
Auto-scaled (2-10x)	$100-500	800-4000 req/s

Recommendation: Use CPU with quantization for cost-effectiveness.