train_stage2.py

"""DPNet supervised learning — two-phase classifier training.

Phase 1 : STDP weight training with teacher signals.
Phase 2 : Memory layer weights are frozen; last-sub-layer spike-time
    representations are extracted and used to train a fully-connected classifier.

Loads a model produced by unsupervised_growth.py.

Example usage:
    python supervised_learning.py -p ./output/mnist/model.net

    # Save model after phase 1, then run phase 2
    python supervised_learning.py -p ./output/mnist/model.net -s --output-dir ./output/mnist

"""

import argparse
import os

import numpy as np
import torch
import torch.nn as nn
import torch.utils.data

from train.my_train import Train
import nest
from nest_interface.interface_base import InterfaceBase



_DATASET_DEFAULTS = {
    'mnist': dict(input_size=28 * 28, n_classes=10, converter='power',
                  output_dir='./output/mnist'),
    'cifar': dict(input_size=32*32*3, n_classes=10, converter='power',
                  output_dir='./output/cifar10'),
}


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser(
        description="DPNet two-phase supervised learning",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # ---- Dataset ----
    p.add_argument("--dataset", type=str, default="mnist", choices=["mnist", "cifar"])
    p.add_argument("--data-dir", type=str, default="./data/data_set/")
    p.add_argument("--train-samples", type=int, default=300,
                   help="Training samples per class")
    p.add_argument("--test-samples", type=int, default=20,
                   help="Test samples per class")

    # ---- Model persistence ----
    p.add_argument("-p", "--path", type=str, required=True,
                   help="Path to a grown .net model file (from unsupervised_growth.py)")
    p.add_argument("-s", "--save", action="store_true", default=False,
                   help="Save the model after phase 1 to --output-dir/model_learned.net")
    p.add_argument("--output-dir", type=str, default=None,
                   help="Output directory (default: ./output/<dataset>)")

    # ---- Network architecture (must match the grown model) ----
    p.add_argument("--memory-size", type=int, nargs=3, default=[100, 10, 2],
                   metavar=("X", "Y", "Z"))
    p.add_argument("--n-classes", type=int, default=None)
    p.add_argument("--mem-synapse-rate", type=float, default=0.0009)
    p.add_argument("--inp-out-degree", type=int, default=40)
    p.add_argument("--inp-max-weight", type=float, default=20.0)

    # ---- Simulator ----
    p.add_argument("--duration", type=float, default=500.0)
    p.add_argument("--t-ref", type=float, default=20.0)
    p.add_argument("--v-th", type=float, default=-68.0)

    # ---- Phase 1: teacher-signal learning ----
    p.add_argument("--skip-nest-phase", action="store_true", default=False,
                   help="Skip phase 1 and go straight to phase 2")
    p.add_argument("--learn-iters", type=int, default=5,
                   help="Number of teacher-signal learning epochs")
    p.add_argument("--max-learning-iters", type=int, default=15,
                   help="Max teacher-signal iterations per sample")
    p.add_argument("--inhibitory-time", type=float, default=110.0,
                   help="Initial teacher inhibitory time (ms)")
    p.add_argument("--excitatory-time", type=float, default=140.0,
                   help="Initial teacher excitatory time (ms)")
    p.add_argument("--fip-threshold", type=int, default=9)
    p.add_argument("--fid-threshold", type=int, default=3)
    p.add_argument("--fip-weight", type=float, default=800.0)

    # ---- Phase 2: classifier ----
    p.add_argument("--epochs", type=int, default=500,
                   help="Training epochs")
    p.add_argument("--lr", type=float, default=1e-4,
                   help="Adam learning rate")
    p.add_argument("--batch-size", type=int, default=100,
                   help="Mini-batch size")
    p.add_argument("--hidden-size", type=int, default=1024,
                   help="Hidden layer width in the classifier")

    return p.parse_args()


# ---------------------------------------------------------------------------
# Dataset loaders
# ---------------------------------------------------------------------------

def load_dataset(args, trainer):
    if args.dataset == 'mnist':
        import data.data_import.read_mnist as reader
        return trainer.select_my_data_set(reader.read_data_sets(args.data_dir))
    else:
        import data.data_import.read_cifar as reader
        return trainer.select_my_data_set(
            reader.read_data_sets(args.data_dir))


# ---------------------------------------------------------------------------
# Phase 1: teacher-signal learning
# ---------------------------------------------------------------------------

def run_teacher_signal_learning(trainer, train_data, train_target, args):
    """Teacher-signal STDP learning on the output layer."""
    trainer.connect_memory2output_after(train_data, train_target, connect_th=0.5, weight=1180)
    trainer.learning(
        train_data, train_target,
        iteration_all=args.learn_iters,
        inhibitory_time=args.inhibitory_time,
        excitatory_time=args.excitatory_time,
    )


# ---------------------------------------------------------------------------
# Phase 2: feature extraction + classifier
# ---------------------------------------------------------------------------

def extract_memory_features(trainer, datas, targets):
    """Run samples through the frozen memory network; return last-sub-layer spike-time features.

    Each neuron in the last memory sub-layer contributes one feature: its earliest
    spike time relative to the simulation window start (0.0 = no spike).

    Returns:
        features : np.ndarray, shape (N, layer_size)
        targets  : np.ndarray, shape (N,)
    """
    trainer.close_memory_stdp()

    one_layer = trainer.memory_size[0] * trainer.memory_size[1]
    n_sublayers = trainer.memory_size[2]
    mem_start = trainer.network.train_layers['memory'][0]
    last_layer_start = mem_start + one_layer * (n_sublayers - 1)
    last_layer_end   = last_layer_start + one_layer
    duration = trainer.interface.duration

    features, feat_targets = [], []
    total = len(datas)
    for idx, (data, target) in enumerate(zip(datas, targets)):
        trainer.network.clear_all_detector()
        trainer.network.load_data(data)
        trainer.interface.simulate()

        senders, spike_times = trainer.interface.get_spikes(trainer.network.detector['memory'])
        post_time = nest.GetKernelStatus()['time']

        feat = np.zeros(one_layer, dtype=np.float32)
        earliest: dict[int, float] = {}
        for s, t in zip(senders, spike_times):
            if last_layer_start <= s < last_layer_end:
                rel_idx  = s - last_layer_start
                rel_time = t - post_time + duration   # normalise to [0, duration]
                if rel_idx not in earliest or rel_time < earliest[rel_idx]:
                    earliest[rel_idx] = rel_time
        for rel_idx, t in earliest.items():
            feat[rel_idx] = t

        features.append(feat)
        feat_targets.append(target)

        if (idx + 1) % 50 == 0 or (idx + 1) == total:
            print(f"  Feature extraction: {idx + 1}/{total}")

    return np.array(features), np.array(feat_targets, dtype=np.int64)


class MemoryClassifier(nn.Module):
    """Fully-connected classifier on top of frozen memory representations."""

    def __init__(self, input_size: int, n_classes: int, hidden_size: int = 1024):
        super().__init__()
        self.net = nn.Sequential(
            nn.BatchNorm1d(input_size),
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.BatchNorm1d(hidden_size),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.BatchNorm1d(hidden_size),
            nn.Linear(hidden_size, n_classes),
        )

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


def run_classifier(train_feats, train_targets, test_feats, test_targets,
                       n_classes, hidden_size, epochs, lr, batch_size, output_dir):
    """Train and evaluate the classifier on extracted memory features."""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"device: {device}")

    X_train = torch.from_numpy(train_feats).to(device)
    y_train = torch.from_numpy(train_targets).to(device)
    X_test  = torch.from_numpy(test_feats).to(device)
    y_test  = torch.from_numpy(test_targets).to(device)

    train_ds = torch.utils.data.TensorDataset(X_train, y_train)
    test_ds  = torch.utils.data.TensorDataset(X_test,  y_test)
    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=batch_size,
                                               shuffle=True, drop_last=True)
    test_loader  = torch.utils.data.DataLoader(test_ds,  batch_size=batch_size)

    model     = MemoryClassifier(X_train.shape[1], n_classes, hidden_size).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    best_test_acc = 0.0
    for epoch in range(1, epochs + 1):
        model.train()
        train_correct = train_total = 0
        for x, y in train_loader:
            optimizer.zero_grad()
            logits = model(x)
            loss = criterion(logits, y)
            loss.backward()
            optimizer.step()
            preds = logits.argmax(dim=1)
            train_correct += (preds == y).sum().item()
            train_total   += y.size(0)

        model.eval()
        test_correct = test_total = 0
        with torch.no_grad():
            for x, y in test_loader:
                preds = model(x).argmax(dim=1)
                test_correct += (preds == y).sum().item()
                test_total   += y.size(0)

        train_acc = 100 * train_correct / train_total
        test_acc  = 100 * test_correct  / test_total
        best_test_acc = max(best_test_acc, test_acc)
        print(f"Epoch {epoch:4d}/{epochs}  "
              f"train {train_acc:.2f}%  test {test_acc:.2f}%")

    print(f"\nBest test accuracy: {best_test_acc:.2f}%")

    # Save model
    os.makedirs(output_dir, exist_ok=True)
    torch_path = os.path.join(output_dir, "classifier.pt")
    torch.save(model.state_dict(), torch_path)
    print(f"Classifier saved to {torch_path}")

    return model


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    args = parse_args()

    if not os.path.isfile(args.path):
        raise FileNotFoundError(
            f"Model file not found: {args.path}\n"
            "Run unsupervised_growth.py -s first to generate a model.")

    ds = _DATASET_DEFAULTS[args.dataset]
    n_classes  = args.n_classes  or ds['n_classes']
    output_dir = args.output_dir or ds['output_dir']

    interface = InterfaceBase(duration=args.duration, t_ref=args.t_ref, V_th=args.v_th)

    trainer = Train(
        interface,
        input_size=int(ds['input_size']),
        memory_size=args.memory_size,
        output_size=int(n_classes),
        fip_threshold=args.fip_threshold,
        fid_threshold=args.fid_threshold,
        fip_weight=args.fip_weight,
        max_learning_iters=args.max_learning_iters,
        output_dir=str(output_dir),
    )

    network = trainer.create_net(
        converter=ds['converter'],
        load=True,
        load_path=args.path,
        mem_synapse_rate=args.mem_synapse_rate,
        inp_out_degree=args.inp_out_degree,
        inp_max_weight=args.inp_max_weight,
    )

    interface.close_multimeter(network.multimeter["input"])
    interface.close_multimeter(network.multimeter["memory"])
    interface.close_multimeter(network.multimeter["output"])

    datas, targets, test_datas, test_targets = load_dataset(args, trainer)

    train_data, train_target = [], []
    test_data,  test_target  = [], []
    for i in range(n_classes):
        train_data.extend(datas[i][:args.train_samples])
        train_target.extend(targets[i][:args.train_samples])
        test_data.extend(test_datas[i][:args.test_samples])
        test_target.extend(test_targets[i][:args.test_samples])

    # ------------------------------------------------------------------
    # Phase 1: teacher-signal learning
    # ------------------------------------------------------------------
    if not args.skip_nest_phase:
        print("\n=== Phase 1: teacher-signal learning ===")
        run_teacher_signal_learning(trainer, train_data, train_target, args)

        if args.save:
            os.makedirs(output_dir, exist_ok=True)
            trainer.save_model(os.path.join(output_dir, "model_learned.net"))

    # ------------------------------------------------------------------
    # Phase 2: extract frozen memory representations to classifier
    # ------------------------------------------------------------------
    print("\n=== Phase 2: extracting memory representations ===")
    trainer.close_memory_stdp()

    print("  Processing training set...")
    train_feats, train_feat_targets = extract_memory_features(
        trainer, train_data, train_target)

    print("  Processing test set...")
    test_feats, test_feat_targets = extract_memory_features(
        trainer, test_data, test_target)

    feat_dim = train_feats.shape[1]
    print(f"  Feature dimension: {feat_dim}  "
          f"(train: {len(train_feats)}, test: {len(test_feats)})")

    print("\n=== Phase 2: training classifier ===")
    run_classifier(
        train_feats, train_feat_targets,
        test_feats,  test_feat_targets,
        n_classes=n_classes,
        hidden_size=args.hidden_size,
        epochs=args.epochs,
        lr=args.lr,
        batch_size=args.batch_size,
        output_dir=output_dir,
    )


if __name__ == "__main__":
    main()