Create train.py

train.py

@@ -0,0 +1,111 @@
+"""
+Training script for Spatial Context Networks (SCN).
+
+Example usage:
+    python train.py --input_dim 10 --n_neurons 32 --output_dim 4 --epochs 50
+"""
+
+import argparse
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, TensorDataset
+
+from spatial_context_networks import SpatialContextNetwork
+
+
+def make_synthetic_dataset(n_samples=256, input_dim=10, output_dim=4, seed=42):
+    """Creates a simple synthetic classification dataset for demonstration."""
+    torch.manual_seed(seed)
+    X = torch.randn(n_samples, input_dim)
+    y = torch.randint(0, output_dim, (n_samples,))
+    return TensorDataset(X, y)
+
+
+def train(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Training on: {device}")
+
+    # Data
+    dataset = make_synthetic_dataset(
+        n_samples=args.n_samples,
+        input_dim=args.input_dim,
+        output_dim=args.output_dim,
+    )
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
+
+    # Model
+    model = SpatialContextNetwork(
+        input_dim=args.input_dim,
+        n_neurons=args.n_neurons,
+        output_dim=args.output_dim,
+        routing_threshold=args.routing_threshold,
+        stability_factor=args.stability_factor,
+        explosion_threshold=args.explosion_threshold,
+    ).to(device)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {total_params}")
+    print(model)
+
+    optimizer = optim.Adam(model.parameters(), lr=args.lr)
+    criterion = nn.CrossEntropyLoss()
+
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        total_loss = 0.0
+        correct = 0
+
+        for X_batch, y_batch in loader:
+            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
+            optimizer.zero_grad()
+            logits = model(X_batch)
+            loss = criterion(logits, y_batch)
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item() * len(X_batch)
+            correct += (logits.argmax(dim=-1) == y_batch).sum().item()
+
+        avg_loss = total_loss / len(dataset)
+        accuracy = correct / len(dataset)
+
+        if epoch % 10 == 0 or epoch == 1:
+            # Network efficiency stats
+            model.eval()
+            with torch.no_grad():
+                sample_x = torch.randn(args.batch_size, args.input_dim).to(device)
+                stats = model.get_network_stats(sample_x)
+            print(
+                f"Epoch {epoch:3d}/{args.epochs} | "
+                f"Loss: {avg_loss:.4f} | Acc: {accuracy:.3f} | "
+                f"Active neurons: {stats['mean_active_neurons']:.1f}/{args.n_neurons} "
+                f"(eff={stats['network_efficiency']:.2f}) | "
+                f"Context score: {stats['mean_context_score']:.3f}"
+            )
+
+    if args.save_path:
+        torch.save(model.state_dict(), args.save_path)
+        print(f"\nModel saved to {args.save_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train a Spatial Context Network")
+
+    # Architecture
+    parser.add_argument("--input_dim", type=int, default=10)
+    parser.add_argument("--n_neurons", type=int, default=32)
+    parser.add_argument("--output_dim", type=int, default=4)
+    parser.add_argument("--routing_threshold", type=float, default=0.5)
+    parser.add_argument("--stability_factor", type=float, default=10.0)
+    parser.add_argument("--explosion_threshold", type=float, default=2.0)
+
+    # Training
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--n_samples", type=int, default=256)
+    parser.add_argument("--save_path", type=str, default=None)
+
+    args = parser.parse_args()
+    train(args)