From adeb11184fa64e7765080bfd2e8ac2297ef66926 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Tue, 14 Apr 2026 09:35:05 +0200 Subject: [PATCH 01/19] Rename STDWI module --- symmnet/src/symmnet/{stdwi_original.py => stdwi.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename symmnet/src/symmnet/{stdwi_original.py => stdwi.py} (100%) diff --git a/symmnet/src/symmnet/stdwi_original.py b/symmnet/src/symmnet/stdwi.py similarity index 100% rename from symmnet/src/symmnet/stdwi_original.py rename to symmnet/src/symmnet/stdwi.py From 623c323958ca4bed4ffe27531cd4ac196ac3ab67 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Tue, 14 Apr 2026 09:40:01 +0200 Subject: [PATCH 02/19] Add/update typing --- scripts/symm_net/load_utils.py | 28 +++++- scripts/symm_net/main_salnet.py | 59 ++++++++---- scripts/symm_net/salnet_symm.py | 17 ++-- symmnet/src/symmnet/conv_net.py | 33 ++++--- symmnet/src/symmnet/datasets.py | 21 ++++- symmnet/src/symmnet/layers.py | 109 +++++++++++++--------- symmnet/src/symmnet/rdd_layers.py | 32 +++++-- symmnet/src/symmnet/rdd_net.py | 29 ++++-- symmnet/src/symmnet/sal_net.py | 147 +++++++++++++++--------------- symmnet/src/symmnet/stdwi.py | 103 +++++++++++++-------- symmnet/src/symmnet/utils.py | 9 +- 11 files changed, 361 insertions(+), 226 deletions(-) diff --git a/scripts/symm_net/load_utils.py b/scripts/symm_net/load_utils.py index 7a9042c..a3b86ac 100644 --- a/scripts/symm_net/load_utils.py +++ b/scripts/symm_net/load_utils.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import argparse +from typing import Any, Callable import yaml @@ -34,14 +35,22 @@ }, "dataset": "cifar10", } -ALLOWED_DATASETS = {"cifar10": cifar10, "mnist": mnist, "fmnist": fmnist, "svhn": svhn} +ALLOWED_DATASETS: dict[str, Callable] = { + "cifar10": cifar10, + "mnist": mnist, + "fmnist": fmnist, + "svhn": svhn, +} -def parse_tags(s): +def parse_tags(s: str | None) -> list[str]: return [tag.strip() for tag in s.split(",")] if s else [] -def load_params(param_file, section=None): +def load_params( + param_file: str, + section: str | None = None, +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any], str]: with open(param_file, "r") as f: content = yaml.safe_load(f) if section: @@ -61,13 +70,22 @@ def load_params(param_file, section=None): ) -def merge(base, override): +def merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: result = base.copy() result.update(override) return result -def settings_loader(): +def settings_loader() -> tuple[ + dict[str, Any], + dict[str, Any], + dict[str, Any], + Callable, + list[str], + list[str], + str | None, + str, +]: parser = argparse.ArgumentParser() parser.add_argument("-f", type=str, help="Path to parameter file.") parser.add_argument("-s", type=str, help="Section name in YAML file.") diff --git a/scripts/symm_net/main_salnet.py b/scripts/symm_net/main_salnet.py index 8578901..59ef170 100644 --- a/scripts/symm_net/main_salnet.py +++ b/scripts/symm_net/main_salnet.py @@ -4,6 +4,7 @@ import os import shutil from datetime import datetime +from typing import Any import matplotlib @@ -11,6 +12,10 @@ import matplotlib.pyplot as plt import numpy as np # noqa import torch +import torch.nn as nn +import torch.optim as optim +from torch import Tensor +from torch.utils.data import DataLoader from tqdm import tqdm from load_utils import settings_loader @@ -85,7 +90,7 @@ # --------------------------- -def create_run_dirs(base_dir="runs", tags=None): +def create_run_dirs(base_dir: str = "runs", tags: list[str] | None = None) -> str: """Create a timestamped run directory with subfolders for figs and checkpoints.""" timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") tag_str = "_".join(tags) if tags else "run" @@ -97,13 +102,13 @@ def create_run_dirs(base_dir="runs", tags=None): return run_dir -def save_json(obj, path): +def save_json(obj: Any, path: str) -> None: """Serialize obj to a JSON file at path.""" with open(path, "w") as f: json.dump(obj, f, indent=2) -def append_metric(metrics_dict, key, value): +def append_metric(metrics_dict: dict[str, Any], key: str, value: float) -> None: """Append a scalar value to a list under key in metrics_dict.""" if key not in metrics_dict: metrics_dict[key] = [] @@ -121,7 +126,7 @@ def append_metric(metrics_dict, key, value): shutil.copy(param_file, os.path.join(run_dir, os.path.basename(param_file))) # in-memory metrics store -metrics = { +metrics: dict[str, Any] = { "params": params, "sal_params": sal_params, "rdd_params": rdd_params, @@ -188,20 +193,27 @@ def append_metric(metrics_dict, key, value): from symmnet import RDDNet, dt, input_rate, mem rdd_net = RDDNet(net.feature_layers_sizes) - symm_losses = [[] for _ in range(3)] - decay_losses = [[] for _ in range(3)] - sparse_losses = [[] for _ in range(3)] - self_losses = [[] for _ in range(3)] - amp_losses = [[] for _ in range(3)] - info_losses = [[] for _ in range(3)] - corr_percents = [[] for _ in range(3)] + symm_losses: list[list[float]] = [[] for _ in range(3)] + decay_losses: list[list[float]] = [[] for _ in range(3)] + sparse_losses: list[list[float]] = [[] for _ in range(3)] + self_losses: list[list[float]] = [[] for _ in range(3)] + amp_losses: list[list[float]] = [[] for _ in range(3)] + info_losses: list[list[float]] = [[] for _ in range(3)] + corr_percents: list[list[float]] = [[] for _ in range(3)] # --------------------------- # Training and Evaluation Functions # --------------------------- -def train(model, optimizer, criterion, train_loader, device, metrics): +def train( + model: nn.Module, + optimizer: optim.Optimizer, + criterion: nn.Module, + train_loader: DataLoader, + device: str, + metrics: dict[str, Any], +) -> None: """Standard training loop for one epoch.""" model.train() train_loss = 0 @@ -232,7 +244,13 @@ def train(model, optimizer, criterion, train_loader, device, metrics): append_metric(metrics["scalars"], "accuracy/train", 100 * correct / total) -def test(model, criterion, test_loader, device, metrics): +def test( + model: nn.Module, + criterion: nn.Module, + test_loader: DataLoader, + device: str, + metrics: dict[str, Any], +) -> None: """Evaluation loop on the test set.""" model.eval() total_loss = 0 @@ -261,7 +279,7 @@ def test(model, criterion, test_loader, device, metrics): append_metric(metrics["scalars"], "accuracy/test", 100 * correct / total) -def train_sal(sal_net, sal_params): +def train_sal(sal_net: SALNet, sal_params: dict[str, Any]) -> list[Tensor]: """Run one SAL phase (spike-based feedback weight update).""" progress_bar = tqdm(range(sal_params["len_epoch"] * sal_params["t_ref"])) with torch.no_grad(): @@ -273,7 +291,12 @@ def train_sal(sal_net, sal_params): return dw -def eval_symmetry(net, metrics, run_dir, epoch): +def eval_symmetry( + net: ConvNet, + metrics: dict[str, Any], + run_dir: str, + epoch: int, +) -> None: """Evaluate forward/feedback weight symmetry; save scatter plots every SCATTER_INTERVAL epochs.""" weights = list(net.parameters_weight()) fb_weights = list(net.parameters_fb_weight(ignore_require_grad=True)) @@ -313,7 +336,7 @@ def eval_symmetry(net, metrics, run_dir, epoch): plt.close(fig) -def save_final_plots(metrics, run_dir): +def save_final_plots(metrics: dict[str, Any], run_dir: str) -> None: """Generate and save final summary plots for loss, accuracy, and symmetry angle.""" scalars = metrics["scalars"] figs_dir = os.path.join(run_dir, "figs") @@ -359,7 +382,7 @@ def save_final_plots(metrics, run_dir): # --------------------------- -def train_rdd(): +def train_rdd() -> None: """Implement the spike-based RDD feed-back learning.""" rdd_net.reset() @@ -560,7 +583,7 @@ def train_rdd(): # --------------------------- -def main(): +def main() -> None: eval_symmetry(net, metrics, run_dir, epoch=0) for epoch in range(params["n_epochs"]): diff --git a/scripts/symm_net/salnet_symm.py b/scripts/symm_net/salnet_symm.py index 493fdc9..30edb18 100644 --- a/scripts/symm_net/salnet_symm.py +++ b/scripts/symm_net/salnet_symm.py @@ -4,6 +4,7 @@ import json import os from datetime import datetime +from typing import Any import matplotlib @@ -41,7 +42,7 @@ # --------------------------- -def parse_tags(s): +def parse_tags(s: str | None) -> list[str]: return [tag.strip() for tag in s.split(",")] if s else [] @@ -70,7 +71,7 @@ def parse_tags(s): # --------------------------- -def create_run_dirs(base_dir="runs", tags=None): +def create_run_dirs(base_dir: str = "runs", tags: list[str] | None = None) -> str: """Create a timestamped run directory with subfolders for figs.""" timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") tag_str = "_".join(tags) if tags else "pure_sal" @@ -81,13 +82,13 @@ def create_run_dirs(base_dir="runs", tags=None): return run_dir -def save_json(obj, path): +def save_json(obj: Any, path: str) -> None: """Serialize obj to a JSON file at path.""" with open(path, "w") as f: json.dump(obj, f, indent=2) -def append_metric(scalars, key, value): +def append_metric(scalars: dict[str, Any], key: str, value: float) -> None: """Append a scalar value to a list under key in scalars.""" if key not in scalars: scalars[key] = [] @@ -100,7 +101,7 @@ def append_metric(scalars, key, value): run_dir = create_run_dirs(base_dir=args.output_dir, tags=tags) -metrics = { +metrics: dict[str, Any] = { "params": params, "scalars": {}, } @@ -125,7 +126,7 @@ def append_metric(scalars, key, value): # --------------------------- -def eval_and_log_symmetry(scalars, epoch): +def eval_and_log_symmetry(scalars: dict[str, Any], epoch: int) -> None: """Log symmetry metrics and save scatter plots every PLOT_INTERVAL epochs.""" for i in range(n_layers - 1): w = net.layers[i + 1].weight.detach() @@ -157,7 +158,7 @@ def eval_and_log_symmetry(scalars, epoch): plt.close(fig) -def save_final_plots(metrics, run_dir): +def save_final_plots(metrics: dict[str, Any], run_dir: str) -> None: """Generate and save final summary plots for symmetry angle and corrcoef.""" scalars = metrics["scalars"] figs_dir = os.path.join(run_dir, "figs") @@ -194,7 +195,7 @@ def save_final_plots(metrics, run_dir): # --------------------------- -def main(): +def main() -> None: # log symmetry before any training eval_and_log_symmetry(metrics["scalars"], epoch=0) save_json(metrics, os.path.join(run_dir, "metrics.json")) diff --git a/symmnet/src/symmnet/conv_net.py b/symmnet/src/symmnet/conv_net.py index 50864ff..82562d8 100755 --- a/symmnet/src/symmnet/conv_net.py +++ b/symmnet/src/symmnet/conv_net.py @@ -1,21 +1,24 @@ """Define the ConvNet architecture (standard pytorch ANN).""" import math +from typing import Generator import torch.nn as nn +from torch import Tensor +from torch.nn import Parameter from .layers import Conv2dFA, LinearFA, LinearKP, LinearSCFA -def conv2d_outsize(insize, kernel_size): +def conv2d_outsize(insize: int, kernel_size: int) -> int: return insize - kernel_size + 1 -def maxpool2d_outsize(insize, kernel_size, stride): +def maxpool2d_outsize(insize: int, kernel_size: int, stride: int) -> int: return math.floor((insize - kernel_size) / 2 + 1) -def feature_layer_outsize(insize): +def feature_layer_outsize(insize: int) -> int: return maxpool2d_outsize( conv2d_outsize(maxpool2d_outsize(conv2d_outsize(insize, 5), 2, 2), 5), 2, 2 ) @@ -38,13 +41,13 @@ class ConvNet(nn.Module): def __init__( self, - input_channels, - input_shape, - use_backprop=False, - use_kp=False, - use_scfa=False, - use_fa_conv_layers=False, - ): + input_channels: int, + input_shape: tuple[int, int], + use_backprop: bool = False, + use_kp: bool = False, + use_scfa: bool = False, + use_fa_conv_layers: bool = False, + ) -> None: """ Initialize the ConvNet model. @@ -115,7 +118,7 @@ def __init__( # Combine all layers into a single sequential module for the forward pass self.out = nn.Sequential(*(self.feature_layers + self.classification_layers)) - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: """ Forward pass of the ConvNet. @@ -127,19 +130,21 @@ def forward(self, x): """ return self.out(x) - def parameters_weight(self): + def parameters_weight(self) -> Generator[Parameter, None, None]: for name, param in self.named_parameters(): if name in ConvNet.FC_WEIGHT: yield param - def parameters_fb_weight(self, ignore_require_grad=False): + def parameters_fb_weight( + self, ignore_require_grad: bool = False + ) -> Generator[Parameter, None, None]: for name, param in self.named_parameters(): if name in ConvNet.FC_FB_WEIGHT and ( param.requires_grad or ignore_require_grad ): yield param - def parameters_other(self): + def parameters_other(self) -> Generator[Parameter, None, None]: for name, param in self.named_parameters(): if name not in ConvNet.FC_WEIGHT + ConvNet.FC_FB_WEIGHT: yield param diff --git a/symmnet/src/symmnet/datasets.py b/symmnet/src/symmnet/datasets.py index d4b3095..df7f8fb 100644 --- a/symmnet/src/symmnet/datasets.py +++ b/symmnet/src/symmnet/datasets.py @@ -1,10 +1,13 @@ #!/usr/bin/env python3 import torch +from torch.utils.data import DataLoader from torchvision import datasets, transforms -def cifar10(batch_size, path="../datasets"): +def cifar10( + batch_size: int, path: str = "../datasets" +) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 3 shape = (32, 32) # Data augmentation and normalization for training @@ -42,7 +45,9 @@ def cifar10(batch_size, path="../datasets"): return train_loader, test_loader, n_channels, shape -def mnist(batch_size, path="../datasets"): +def mnist( + batch_size: int, path: str = "../datasets" +) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 1 shape = (28, 28) # Data augmentation and normalization for training @@ -80,7 +85,9 @@ def mnist(batch_size, path="../datasets"): return train_loader, test_loader, n_channels, shape -def fmnist(batch_size, path="../datasets"): +def fmnist( + batch_size: int, path: str = "../datasets" +) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 1 shape = (28, 28) # Data augmentation and normalization for training @@ -118,7 +125,9 @@ def fmnist(batch_size, path="../datasets"): return train_loader, test_loader, n_channels, shape -def imagenette(batch_size, path="../datasets"): +def imagenette( + batch_size: int, path: str = "../datasets" +) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 3 shape = (160, 160) # Data augmentation and normalization for training @@ -156,7 +165,9 @@ def imagenette(batch_size, path="../datasets"): return train_loader, test_loader, n_channels, shape -def svhn(batch_size, path="../datasets"): +def svhn( + batch_size: int, path: str = "../datasets" +) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 3 shape = (32, 32) # Data augmentation and normalization for training diff --git a/symmnet/src/symmnet/layers.py b/symmnet/src/symmnet/layers.py index c73bcf7..9fc50aa 100755 --- a/symmnet/src/symmnet/layers.py +++ b/symmnet/src/symmnet/layers.py @@ -5,12 +5,17 @@ import torch import torch.nn as nn import torch.nn.functional as F +from torch import Tensor from torch.nn import init from torch.nn.modules.utils import _pair from torch.nn.parameter import Parameter -def conv2d_fa_backward_hook(module, grad_input, grad_output): +def conv2d_fa_backward_hook( + module: nn.Module, + grad_input: tuple[Tensor, ...], + grad_output: tuple[Tensor, ...], +) -> tuple[Tensor, ...] | None: """ Backward hook for Conv2dFA layers. @@ -41,12 +46,16 @@ def conv2d_fa_backward_hook(module, grad_input, grad_output): class LinearFuncFA(torch.autograd.Function): @staticmethod - def forward(ctx, inp, weight, bias, fb_weight): + def forward( + ctx, inp: Tensor, weight: Tensor, bias: Tensor | None, fb_weight: Tensor + ) -> Tensor: ctx.save_for_backward(inp, weight, bias, fb_weight) return F.linear(inp, weight, bias) @staticmethod - def backward(ctx, grad_output): + def backward( + ctx, grad_output: Tensor + ) -> tuple[Tensor, Tensor, Tensor | None, None]: inp, weight, bias, fb_weight = ctx.saved_tensors grad_input = grad_output.mm(fb_weight) @@ -60,12 +69,16 @@ def backward(ctx, grad_output): class LinearFuncKP(torch.autograd.Function): @staticmethod - def forward(ctx, inp, weight, bias, fb_weight): + def forward( + ctx, inp: Tensor, weight: Tensor, bias: Tensor | None, fb_weight: Tensor + ) -> Tensor: ctx.save_for_backward(inp, weight, bias, fb_weight) return F.linear(inp, weight, bias) @staticmethod - def backward(ctx, grad_output): + def backward( + ctx, grad_output: Tensor + ) -> tuple[Tensor, Tensor, Tensor | None, Tensor]: inp, weight, bias, fb_weight = ctx.saved_tensors grad_input = grad_output.mm(fb_weight) @@ -80,12 +93,16 @@ def backward(ctx, grad_output): class LinearFuncSCFA(torch.autograd.Function): @staticmethod - def forward(ctx, inp, weight, bias, fb_weight): + def forward( + ctx, inp: Tensor, weight: Tensor, bias: Tensor | None, fb_weight: Tensor + ) -> Tensor: ctx.save_for_backward(inp, weight, bias, fb_weight) return F.linear(inp, weight, bias) @staticmethod - def backward(ctx, grad_output): + def backward( + ctx, grad_output: Tensor + ) -> tuple[Tensor, Tensor, Tensor | None, None]: inp, weight, bias, fb_weight = ctx.saved_tensors grad_input = grad_output.mm(fb_weight * weight.sign()) @@ -97,15 +114,21 @@ def backward(ctx, grad_output): return grad_input, grad_weight, grad_bias, None -def linear_fa(inp, weight, bias, fb_weights): +def linear_fa( + inp: Tensor, weight: Tensor, bias: Tensor | None, fb_weights: Tensor +) -> Tensor: return LinearFuncFA.apply(inp, weight, bias, fb_weights) -def linear_kp(inp, weight, bias, fb_weights): +def linear_kp( + inp: Tensor, weight: Tensor, bias: Tensor | None, fb_weights: Tensor +) -> Tensor: return LinearFuncKP.apply(inp, weight, bias, fb_weights) -def linear_scfa(inp, weight, bias, fb_weights): +def linear_scfa( + inp: Tensor, weight: Tensor, bias: Tensor | None, fb_weights: Tensor +) -> Tensor: return LinearFuncSCFA.apply(inp, weight, bias, fb_weights) @@ -125,7 +148,7 @@ class LinearFA(nn.Module): __constants__ = ["bias", "in_features", "out_features"] - def __init__(self, in_features, out_features, bias=True): + def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None: """ Initialize the LinearFA layer. @@ -147,7 +170,7 @@ def __init__(self, in_features, out_features, bias=True): self.register_parameter("bias", None) self.reset_parameters() - def reset_parameters(self): + def reset_parameters(self) -> None: """ Initialize or reset the parameters of the layer. """ @@ -158,7 +181,7 @@ def reset_parameters(self): bound = 1 / math.sqrt(fan_in) init.uniform_(self.bias, -bound, bound) - def forward(self, input): + def forward(self, input: Tensor) -> Tensor: """ Forward pass of the LinearFA layer. @@ -170,7 +193,7 @@ def forward(self, input): """ return linear_fa(input, self.weight, self.bias, self.fb_weight) - def extra_repr(self): + def extra_repr(self) -> str: """ Extra representation of the module for printing. @@ -183,11 +206,11 @@ def extra_repr(self): class LinearKP(LinearFA): - def __init__(self, in_features, out_features, bias=True): + def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None: super().__init__(in_features, out_features, bias=bias) self.fb_weight.requires_grad_() - def forward(self, input): + def forward(self, input: Tensor) -> Tensor: """ Forward pass of the LinearFA layer. @@ -201,11 +224,11 @@ def forward(self, input): class LinearSCFA(LinearFA): - def __init__(self, in_features, out_features, bias=True): + def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None: super().__init__(in_features, out_features, bias=bias) self.fb_weight.data = self.fb_weight.data.abs() - def forward(self, input): + def forward(self, input: Tensor) -> Tensor: """ Forward pass of the LinearFA layer. @@ -244,18 +267,18 @@ class _ConvNdFA(nn.Module): def __init__( self, - in_channels, - out_channels, - kernel_size, - stride, - padding, - dilation, - transposed, - output_padding, - groups, - bias, - padding_mode, - ): + in_channels: int, + out_channels: int, + kernel_size: tuple[int, ...], + stride: tuple[int, ...], + padding: tuple[int, ...], + dilation: tuple[int, ...], + transposed: bool, + output_padding: tuple[int, ...], + groups: int, + bias: bool, + padding_mode: str, + ) -> None: """ Initialize the _ConvNdFA layer. @@ -314,7 +337,7 @@ def __init__( self.register_parameter("bias", None) self.reset_parameters() - def reset_parameters(self): + def reset_parameters(self) -> None: """ Initialize or reset the parameters of the layer. """ @@ -325,7 +348,7 @@ def reset_parameters(self): bound = 1 / math.sqrt(fan_in) init.uniform_(self.bias, -bound, bound) - def extra_repr(self): + def extra_repr(self) -> str: """ Initialize or reset the parameters of the layer. """ @@ -368,16 +391,16 @@ class Conv2dFA(_ConvNdFA): def __init__( self, - in_channels, - out_channels, - kernel_size, - stride=1, - padding=0, - dilation=1, - groups=1, - bias=True, - padding_mode="zeros", - ): + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int], + stride: int | tuple[int, int] = 1, + padding: int | tuple[int, int] = 0, + dilation: int | tuple[int, int] = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + ) -> None: """ Initialize the Conv2dFA layer. @@ -413,7 +436,7 @@ def __init__( self.register_backward_hook(conv2d_fa_backward_hook) - def forward(self, input): + def forward(self, input: Tensor) -> Tensor: """ Forward pass of the Conv2dFA layer. diff --git a/symmnet/src/symmnet/rdd_layers.py b/symmnet/src/symmnet/rdd_layers.py index eebb086..93d3d4d 100755 --- a/symmnet/src/symmnet/rdd_layers.py +++ b/symmnet/src/symmnet/rdd_layers.py @@ -1,6 +1,7 @@ """Implementation of the dynamics of RDD layers.""" import numpy as np +import numpy.typing as npt from symmnet import ( RDD_eta, @@ -20,7 +21,7 @@ ) -def kappa(x): +def kappa(x: float | int) -> float: """ Computes the difference of exponentials kernel for synaptic current. @@ -35,7 +36,7 @@ def kappa(x): ) -def get_kappas(n=mem): +def get_kappas(n: int = mem) -> npt.NDArray: """ Computes the kappa kernel for the last n time steps. @@ -65,7 +66,12 @@ class SpikingFA: b_input_size (int, optional): Size of the feedback input (enables RDD). """ - def __init__(self, size, f_input_size=None, b_input_size=None): + def __init__( + self, + size: int, + f_input_size: int | None = None, + b_input_size: int | None = None, + ) -> None: self.size = size self.f_input_size = f_input_size self.b_input_size = b_input_size @@ -87,7 +93,7 @@ def __init__(self, size, f_input_size=None, b_input_size=None): self.reset() - def reset(self): + def reset(self) -> None: """ Resets all state variables for the spiking neuron population. """ @@ -110,7 +116,12 @@ def reset(self): self.n_spikes = np.zeros((self.size, 1)) self.max_u = np.zeros((self.size, 1)) - def set_weights(self, weight=None, bias=None, fb_weight=None): + def set_weights( + self, + weight: npt.NDArray | None = None, + bias: npt.NDArray | None = None, + fb_weight: npt.NDArray | None = None, + ) -> None: """ Sets the feedforward and feedback weights, with normalization. @@ -129,7 +140,12 @@ def set_weights(self, weight=None, bias=None, fb_weight=None): if self.fb_weight_std is None: self.fb_weight_std = np.std(self.fb_weight) - def update(self, f_input=None, b_input=None, driving_input=None): + def update( + self, + f_input: npt.NDArray | None = None, + b_input: npt.NDArray | None = None, + driving_input: npt.NDArray | None = None, + ) -> None: """ Advances the state of the neuron population by one time step. @@ -231,7 +247,7 @@ def update(self, f_input=None, b_input=None, driving_input=None): # update spike histories self.spike_hist = np.concatenate([self.spike_hist[:, 1:], self.fired], axis=1) - def update_RDD_estimate(self): + def update_RDD_estimate(self) -> None: """ Updates the RDD regression parameter estimates for neurons at the end of their RDD window. """ @@ -288,7 +304,7 @@ def update_RDD_estimate(self): self.R[end_mask] = 0 - def update_fb_weights(self): + def update_fb_weights(self) -> None: """ Updates the feedback weights based on the current RDD beta values, if nonzero. """ diff --git a/symmnet/src/symmnet/rdd_net.py b/symmnet/src/symmnet/rdd_net.py index 712e542..0d7f267 100755 --- a/symmnet/src/symmnet/rdd_net.py +++ b/symmnet/src/symmnet/rdd_net.py @@ -1,7 +1,9 @@ """Define the architecture of the spiking RDD net.""" import numpy as np +import numpy.typing as npt import torch +import torch.nn as nn from .rdd_layers import SpikingFA @@ -9,7 +11,7 @@ class RDDNetBase: """This is new!""" - def __init__(self, layer_dims): + def __init__(self, layer_dims: list[int]) -> None: self.n_layers = len(layer_dims) assert self.n_layers >= 2, "The network needs at least two layers!" self.classification_layers = [] @@ -21,7 +23,7 @@ def __init__(self, layer_dims): ) self.classification_layers.append(SpikingFA(layer_dims[-1], layer_dims[-2])) - def out(self, *args): + def out(self, *args: npt.NDArray) -> None: """args must be the driving_spike_hist""" assert len(args) == len(self.classification_layers) - 1 @@ -43,14 +45,14 @@ def out(self, *args): # last Layer: receives feedforward from penultimate layer, no feedback, no external drive self.classification_layers[-1].update(self.classification_layers[-2].spike_hist) - def reset(self): + def reset(self) -> None: """ Resets the state of all SpikingFA layers (membrane potentials, spike history, etc.). """ for layer in self.classification_layers: layer.reset() - def update_fb_weights(self): + def update_fb_weights(self) -> None: """ Updates the feedback weights in all but the last SpikingFA layer. """ @@ -73,7 +75,7 @@ class RDDNet: TODO: make this inherit from RDDNetBase! """ - def __init__(self, layer_dims): + def __init__(self, layer_dims: list[int]) -> None: """ Initializes the RDDNet with a fixed architecture of SpikingFA layers. """ @@ -88,7 +90,7 @@ def __init__(self, layer_dims): ) self.classification_layers.append(SpikingFA(layer_dims[3], layer_dims[2])) - def copy_weights_from(self, layers): + def copy_weights_from(self, layers: list[nn.Module]) -> None: """ Copies weights and feedback weights from the linear layers of the corresponting pytorch ANN. @@ -116,7 +118,9 @@ def copy_weights_from(self, layers): layers[4].bias.detach().cpu().numpy().astype(np.float32)[:, np.newaxis], ) - def copy_weights_to(self, layers, device): + def copy_weights_to( + self, layers: list[nn.Module], device: torch.device | str + ) -> None: """ Copies feedback weights from the SpikingFA layers back to the corresponding PyTorch layers, typically after RDD-based updates. @@ -136,7 +140,12 @@ def copy_weights_to(self, layers, device): self.classification_layers[2].fb_weight.astype(np.float32).T ).to(device) - def out(self, driving_spike_hist_1, driving_spike_hist_2, driving_spike_hist_3): + def out( + self, + driving_spike_hist_1: npt.NDArray | None, + driving_spike_hist_2: npt.NDArray | None, + driving_spike_hist_3: npt.NDArray | None, + ) -> None: """ Sequentially updates each SpikingFA layer given the driving spike histories and the spike histories of adjacent layers. @@ -167,14 +176,14 @@ def out(self, driving_spike_hist_1, driving_spike_hist_2, driving_spike_hist_3): # Layer 3: receives feedforward from layer 2, no feedback, no external drive self.classification_layers[3].update(self.classification_layers[2].spike_hist) - def reset(self): + def reset(self) -> None: """ Resets the state of all SpikingFA layers (membrane potentials, spike history, etc.). """ for layer in self.classification_layers: layer.reset() - def update_fb_weights(self): + def update_fb_weights(self) -> None: """ Updates the feedback weights in all but the last SpikingFA layer. """ diff --git a/symmnet/src/symmnet/sal_net.py b/symmnet/src/symmnet/sal_net.py index e501843..b850e33 100644 --- a/symmnet/src/symmnet/sal_net.py +++ b/symmnet/src/symmnet/sal_net.py @@ -5,6 +5,7 @@ import torch import torch.nn as nn +from torch import Tensor from .utils import batched_outer @@ -24,7 +25,7 @@ def __init__( num_neurons: int, buffer_length: int, batch_size: int = 1, - ): + ) -> None: super().__init__() self.buffer_length = buffer_length self.register_buffer( @@ -32,29 +33,29 @@ def __init__( ) @abstractmethod - def append(self, spikes): + def append(self, spikes: Tensor) -> None: pass @abstractmethod - def get(self): + def get(self) -> Tensor: pass class SimpleSpikeBuffer(BaseSpikeBuffer): """The simplest buffer one could think of.""" - def append(self, spikes): + def append(self, spikes: Tensor) -> None: self.buffer = torch.roll(self.buffer, shifts=1, dims=2) self.buffer[:, :, 0] = spikes - def get(self): + def get(self) -> Tensor: return self.buffer class BaseKernel(nn.Module): """Base kernel, both for PSP and STDP.""" - def __init__(self, kernel): + def __init__(self, kernel: Tensor) -> None: super().__init__() assert kernel.shape[0] == 1, "Dimension 0 of the Kernel has to be 1!" assert kernel.shape[1] == 1, "Dimension 1 of the Kernel has to be 1!" @@ -62,7 +63,7 @@ def __init__(self, kernel): self.register_buffer("kernel", kernel) self.len_kernel = self.kernel.shape[2] - def forward(self, spikes): + def forward(self, spikes: Tensor) -> Tensor: # spikes: [batch_size, num_neurons, time_steps] # Use 1D convolution to apply the box filter along the time axis # Reshape for conv1d: [batch_size * num_neurons, 1, time_steps] @@ -78,7 +79,7 @@ def forward(self, spikes): class RectangularPSP(BaseKernel): """Rectangular PSPs.""" - def __init__(self, t_syn): + def __init__(self, t_syn: int) -> None: kernel = torch.ones(1, 1, t_syn) super().__init__(kernel) @@ -86,7 +87,7 @@ def __init__(self, t_syn): class AlphaPSP(BaseKernel): """Alpha-shaped PSP kernel, normalized to t_ref.""" - def __init__(self, t_syn, t_ref): + def __init__(self, t_syn: float, t_ref: float) -> None: ts = torch.arange(int(t_syn * 10)).reshape(1, 1, -1) kernel = t_ref / t_syn**2 * ts * torch.exp(-ts / t_syn) super().__init__(kernel) @@ -95,33 +96,33 @@ def __init__(self, t_syn, t_ref): class ExpSTDP(BaseKernel): """Exponential STDP window.""" - def __init__(self, tau, a, len_kernel): + def __init__(self, tau: float, a: float, len_kernel: int) -> None: ts = torch.arange(len_kernel).reshape(1, 1, -1) kernel = a * torch.exp(-ts / tau) super().__init__(kernel) class ActivationFunction(nn.Module): - def __init__(self, t_ref): + def __init__(self, t_ref: float) -> None: super().__init__() self.log_t_ref = math.log(t_ref) - def forward(self, mem_pot): + def forward(self, mem_pot: Tensor) -> Tensor: return torch.sigmoid(mem_pot - self.log_t_ref) class GLMHiddenLayer(nn.Module): def __init__( self, - n_in, - n_out, - n_next, - t_ref, - stdp_lr, - batch_size=1, - buffer_length=2, - psp=None, - ): + n_in: int, + n_out: int, + n_next: int, + t_ref: int, + stdp_lr: float, + batch_size: int = 1, + buffer_length: int = 2, + psp: BaseKernel | None = None, + ) -> None: super().__init__() # model settings: @@ -157,17 +158,17 @@ def __init__( self.batch_size = batch_size @property - def device(self): + def device(self) -> torch.device: return next(self.parameters()).device - def init_params(self): + def init_params(self) -> None: # initialize params (although they're actually copied from the corresponding ANN) nn.init.kaiming_uniform_(self.weight, a=5**0.5) nn.init.kaiming_uniform_(self.fb_weight, a=5**0.5) bound = self.n_out**-0.5 # TODO check if this is useful nn.init.uniform_(self.bias, -bound, bound) - def update_mempot(self, bottom_up_spikes, top_down_spikes): + def update_mempot(self, bottom_up_spikes: Tensor, top_down_spikes: Tensor) -> None: bottom_up_psps = self.psp(bottom_up_spikes) top_down_psps = self.psp(top_down_spikes) self.mem_pot = ( @@ -176,7 +177,7 @@ def update_mempot(self, bottom_up_spikes, top_down_spikes): + self.bias ) - def update_spikes(self, t=0.0, start_id=1): + def update_spikes(self, t: float = 0.0, start_id: int = 1) -> None: inst_rate = self.phi(self.mem_pot) random_vals = torch.rand(self.batch_size, self.n_out, device=self.device) new_spikes = torch.logical_and( @@ -186,11 +187,11 @@ def update_spikes(self, t=0.0, start_id=1): self.last_spike_counter[new_spikes] = 1.0 self.spikes.append(new_spikes.float()) - def forward(self): + def forward(self) -> None: # not sure, what's the best implementation here... pass - def fb_stdp_online(self, top_down_spikes): + def fb_stdp_online(self, top_down_spikes: Tensor) -> None: # here: top_down_spikes = pre, own spikes = post trace_pre = self.causal_stdp(top_down_spikes) trace_post = self.anticausal_stdp(self.spikes.get()) @@ -202,7 +203,7 @@ def fb_stdp_online(self, top_down_spikes): trace_post[:, :, 0], top_down_spikes[:, :, 0] ) - def apply_fb_weight_update(self, zero_dw=True): + def apply_fb_weight_update(self, zero_dw: bool = True) -> Tensor: dw = self.delta_fb_weight.mean(dim=0) dw /= float(self.causal_stdp.len_kernel) with torch.no_grad(): @@ -215,14 +216,14 @@ def apply_fb_weight_update(self, zero_dw=True): class GLMInputLayer(GLMHiddenLayer): def __init__( self, - n_out, - n_next, - t_ref, - stdp_lr, - batch_size=1, - buffer_length=2, - psp=None, - ): + n_out: int, + n_next: int, + t_ref: int, + stdp_lr: float, + batch_size: int = 1, + buffer_length: int = 2, + psp: BaseKernel | None = None, + ) -> None: super().__init__( 0, n_out, @@ -235,13 +236,13 @@ def __init__( ) self.weight = None - def init_params(self): + def init_params(self) -> None: # initialize params (although they're actually copied from the corresponding ANN) nn.init.kaiming_uniform_(self.fb_weight, a=5**0.5) bound = self.n_out**-0.5 # TODO check if this is useful nn.init.uniform_(self.bias, -bound, bound) - def update_mempot(self, top_down_spikes): + def update_mempot(self, top_down_spikes: Tensor) -> None: top_down_psps = self.psp(top_down_spikes) self.mem_pot = ( torch.matmul(top_down_psps[:, :, 0], self.fb_weight.t()) + self.bias @@ -251,13 +252,13 @@ def update_mempot(self, top_down_spikes): class GLMOutputLayer(GLMHiddenLayer): def __init__( self, - n_in, - n_out, - t_ref, - batch_size=1, - buffer_length=2, - psp=None, - ): + n_in: int, + n_out: int, + t_ref: int, + batch_size: int = 1, + buffer_length: int = 2, + psp: BaseKernel | None = None, + ) -> None: super().__init__( n_in, n_out, @@ -272,24 +273,24 @@ def __init__( self.delta_fb_weight = None self.stdp_lr = None - def init_params(self): + def init_params(self) -> None: # initialize params (although they're actually copied from the corresponding ANN) nn.init.kaiming_uniform_(self.weight, a=5**0.5) bound = self.n_out**-0.5 # TODO check if this is useful nn.init.uniform_(self.bias, -bound, bound) - def update_mempot(self, bottom_up_spikes): + def update_mempot(self, bottom_up_spikes: Tensor) -> None: bottom_up_psps = self.psp(bottom_up_spikes) self.mem_pot = ( torch.matmul(bottom_up_psps[:, :, 0], self.weight.t()) + self.bias ) - def fb_stdp_online(self, bottom_up_spikes): + def fb_stdp_online(self, bottom_up_spikes: Tensor) -> None: raise NotImplementedError( "GLMOutputLayer doesn't have feedback weights and hence no stdp rule for it." ) - def apply_fb_weight_update(self): + def apply_fb_weight_update(self) -> Tensor: raise NotImplementedError( "GLMOutputLayer doesn't have feedback weights and hence no stdp rule for it." ) @@ -298,13 +299,13 @@ def apply_fb_weight_update(self): class SALNetBase(nn.Module): def __init__( self, - layer_dims, - t_ref, - stdp_lr, - batch_size=1, - buffer_length=2, - psp=None, - ): + layer_dims: list[int], + t_ref: int, + stdp_lr: float | list[float], + batch_size: int = 1, + buffer_length: int = 2, + psp: BaseKernel | None = None, + ) -> None: super().__init__() self.n_layers = len(layer_dims) assert self.n_layers >= 2, "The network needs at least two layers!" @@ -353,7 +354,7 @@ def __init__( ) self.layers = nn.ModuleList(layers) - def update_mempot(self): + def update_mempot(self) -> None: # get spikes spikes = [layer.spikes.get() for layer in self.layers] # update layers @@ -362,22 +363,22 @@ def update_mempot(self): self.layers[i].update_mempot(spikes[i - 1], spikes[i + 1]) self.layers[-1].update_mempot(spikes[-2]) - def update_spikes(self, t=0.0): + def update_spikes(self, t: float = 0.0) -> None: for layer in self.layers: layer.update_spikes() - def fb_stdp_online(self): + def fb_stdp_online(self) -> None: for i in range(self.n_layers - 1): self.layers[i].fb_stdp_online(self.layers[i + 1].spikes.get()) - def apply_fb_weight_update(self): - dws = [] # TODO: for debugging + def apply_fb_weight_update(self) -> list[Tensor]: + dws: list[Tensor] = [] # TODO: for debugging for layer in self.layers[:-1]: dw = layer.apply_fb_weight_update() dws.append(dw) return dws - def set_stdp_lr(self, stdp_lr): + def set_stdp_lr(self, stdp_lr: float | list[float]) -> None: if isinstance(stdp_lr, float): self.stdp_lr = [stdp_lr] * 3 for layer, lr in zip(self.layers[:-1], self.stdp_lr): @@ -385,7 +386,7 @@ def set_stdp_lr(self, stdp_lr): class SALNet(SALNetBase): - _COMMON_LAYERS = [ + _COMMON_LAYERS: list[tuple[str, str]] = [ ("fc1.weight", "layers.1.weight"), ("fc1.fb_weight", "layers.0.fb_weight"), ("fc1.bias", "layers.1.bias"), @@ -399,13 +400,13 @@ class SALNet(SALNetBase): def __init__( self, - layer_dims, - t_ref, - stdp_lr, - batch_size=1, - buffer_length=2, - psp=None, - ): + layer_dims: list[int], + t_ref: int, + stdp_lr: float | list[float], + batch_size: int = 1, + buffer_length: int = 2, + psp: BaseKernel | None = None, + ) -> None: assert ( len(layer_dims) == 4 ), "SALNT needs 4 layers to be compatible with the corresponding ConvNet." @@ -418,7 +419,7 @@ def __init__( psp=psp, ) - def load_common_state_dict(self, conv_net_state_dict): + def load_common_state_dict(self, conv_net_state_dict: dict[str, Tensor]) -> None: own_state_dict = self.state_dict() for other, own in SALNet._COMMON_LAYERS: # transpose the feedback weights, because I erroneously definedy the @@ -429,8 +430,8 @@ def load_common_state_dict(self, conv_net_state_dict): own_state_dict[own] = conv_net_state_dict[other] self.load_state_dict(own_state_dict, strict=False) - def get_common_state_dict(self): - other_state_dict = {} + def get_common_state_dict(self) -> dict[str, Tensor]: + other_state_dict: dict[str, Tensor] = {} own_state_dict = self.state_dict() for other, own in SALNet._COMMON_LAYERS: if "fb_weight" in other: diff --git a/symmnet/src/symmnet/stdwi.py b/symmnet/src/symmnet/stdwi.py index b46516c..43b6a7d 100644 --- a/symmnet/src/symmnet/stdwi.py +++ b/symmnet/src/symmnet/stdwi.py @@ -3,6 +3,7 @@ """Copy and simplify code form original repo.""" import numpy as np +import numpy.typing as npt from tqdm import tqdm, trange #################################### @@ -17,7 +18,7 @@ def correlated_poisson_spike_train( simulation_time: float, timestep: float, seed: int = 42, -): +) -> list[npt.NDArray]: """Produces a set of Poisson process sampled spikes with a thinning process to achieve correlation This function creates the desired firing rate as a threshold and draws random numbers (tested against this threshold) to determine spikes. @@ -57,8 +58,13 @@ def correlated_poisson_spike_train( def random_sample_spike_train( - spike_trains, simulation_time, timestep, resample_period, ratio_active, seed=42 -): + spike_trains: list[npt.NDArray], + simulation_time: float, + timestep: float, + resample_period: float, + ratio_active: float, + seed: int = 42, +) -> list[npt.NDArray]: """Randomly samples units from a spike train to be active/inactive. This shifts all spike trains when inactive. Note, this function expects spike trains to have positive only values. @@ -98,7 +104,13 @@ def random_sample_spike_train( return spike_trains -def xpsp_filterer(train, nb_timesteps, timestep, tau_slow, tau_fast): +def xpsp_filterer( + train: npt.NDArray, + nb_timesteps: int, + timestep: float, + tau_slow: float, + tau_fast: float, +) -> npt.NDArray: """Convolves a spike train with a double exponential causal XPSP filter Args: @@ -127,8 +139,12 @@ def xpsp_filterer(train, nb_timesteps, timestep, tau_slow, tau_fast): def spike_trains_to_xpsps( - spike_trains, sim_time, timestep, tau_slow=10.0, tau_fast=3.0 -): + spike_trains: list[npt.NDArray], + sim_time: float, + timestep: float, + tau_slow: float = 10.0, + tau_fast: float = 3.0, +) -> npt.NDArray: """Converts a list of spike trains into a 2D numpy array of post-synaptic potentials Assumes that all spikes cause an equivalent shaped fast E/I PSP -- all psps are positive @@ -153,16 +169,16 @@ def spike_trains_to_xpsps( def lif_dynamics( - xpsps, - weight_matrix, - timestep, - tau=20.0, - thresh=1.0, - rest=0.0, - reset=-1.0, - drift=0.0, - coupling_ratio=1.0, -): + xpsps: npt.NDArray, + weight_matrix: npt.NDArray, + timestep: float, + tau: float = 20.0, + thresh: float = 1.0, + rest: float = 0.0, + reset: float = -1.0, + drift: float = 0.0, + coupling_ratio: float = 1.0, +) -> list[npt.NDArray]: """Computing leaky integrator spiking neuron dynamics given an incident XPSP and weight matrix Args: @@ -214,7 +230,11 @@ def lif_dynamics( return spike_times -def binary_spike_matrix(spike_trains, sim_time, timestep): +def binary_spike_matrix( + spike_trains: list[npt.NDArray], + sim_time: float, + timestep: float, +) -> npt.NDArray: """Converts a list of spike trains into a large NxT binary spike matrix Args: @@ -241,14 +261,14 @@ def binary_spike_matrix(spike_trains, sim_time, timestep): def stdwi_method( - guess_matrix, - input_binary_spikes, - output_binary_spikes, - slow_in_trace, - fast_in_trace, - learning_rate, - decay_weighting, -): + guess_matrix: npt.NDArray, + input_binary_spikes: npt.NDArray, + output_binary_spikes: npt.NDArray, + slow_in_trace: npt.NDArray, + fast_in_trace: npt.NDArray, + learning_rate: float, + decay_weighting: float, +) -> npt.NDArray: """Spike Timing-Dependent based inference of weights Args: @@ -289,7 +309,14 @@ def stdwi_method( return update_matrix -def create_stdwi_trace(prev_trace, binary_spike_matrix, alpha, tau, timestep, alltoall): +def create_stdwi_trace( + prev_trace: npt.NDArray, + binary_spike_matrix: npt.NDArray, + alpha: float, + tau: float, + timestep: float, + alltoall: bool, +) -> npt.NDArray: """Produces an exponential moving average estimation of firing rate from spike times Args: @@ -319,18 +346,18 @@ def create_stdwi_trace(prev_trace, binary_spike_matrix, alpha, tau, timestep, al def apply_stdwi( - inp_spks, - out_spks, - fb_weight, - tau_fast, - tau_slow, - sim_dur, - stim_dur, - lr, - dt, - decay_weighting=0.1, - a_fast=1.0, -): + inp_spks: list[npt.NDArray], + out_spks: list[npt.NDArray], + fb_weight: npt.NDArray, + tau_fast: float, + tau_slow: float, + sim_dur: float, + stim_dur: float, + lr: float, + dt: float, + decay_weighting: float = 0.1, + a_fast: float = 1.0, +) -> npt.NDArray: """My own leaner reimplementation of fitter.stdwi""" n_in = fb_weight.shape[1] diff --git a/symmnet/src/symmnet/utils.py b/symmnet/src/symmnet/utils.py index 3752912..fd821e3 100644 --- a/symmnet/src/symmnet/utils.py +++ b/symmnet/src/symmnet/utils.py @@ -4,17 +4,18 @@ import math import torch +from torch import Tensor -def batched_outer(a, b): +def batched_outer(a: Tensor, b: Tensor) -> Tensor: return a.unsqueeze(2) * b.unsqueeze(1) -def asym_var(a, b): +def asym_var(a: Tensor, b: Tensor) -> Tensor: return (a - b).var() -def asym_angle(a, b): +def asym_angle(a: Tensor, b: Tensor) -> Tensor: assert a.shape == b.shape aa = a.flatten() bb = b.flatten() @@ -22,6 +23,6 @@ def asym_angle(a, b): return torch.acos(torch.clamp(cos, -1.0, 1.0)) * 180.0 / math.pi -def corrcoef(a, b): +def corrcoef(a: Tensor, b: Tensor) -> Tensor: assert a.shape == b.shape return torch.corrcoef(torch.stack([a.flatten(), b.flatten()]))[0, 1] From bd7014abc73c3d7417a7b310203b18efc1b1e5c2 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Tue, 14 Apr 2026 16:15:45 +0200 Subject: [PATCH 03/19] BUGFIX: correct datasets-path --- symmnet/src/symmnet/datasets.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/symmnet/src/symmnet/datasets.py b/symmnet/src/symmnet/datasets.py index df7f8fb..80f4c0d 100644 --- a/symmnet/src/symmnet/datasets.py +++ b/symmnet/src/symmnet/datasets.py @@ -1,12 +1,17 @@ #!/usr/bin/env python3 +from pathlib import Path + import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms +_REPO_ROOT = Path(__file__).parents[3] +_DEFAULT_DATA_DIR = str(_REPO_ROOT / "datasets") + def cifar10( - batch_size: int, path: str = "../datasets" + batch_size: int, path: str = _DEFAULT_DATA_DIR ) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 3 shape = (32, 32) @@ -46,7 +51,7 @@ def cifar10( def mnist( - batch_size: int, path: str = "../datasets" + batch_size: int, path: str = _DEFAULT_DATA_DIR ) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 1 shape = (28, 28) @@ -86,7 +91,7 @@ def mnist( def fmnist( - batch_size: int, path: str = "../datasets" + batch_size: int, path: str = _DEFAULT_DATA_DIR ) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 1 shape = (28, 28) @@ -126,7 +131,7 @@ def fmnist( def imagenette( - batch_size: int, path: str = "../datasets" + batch_size: int, path: str = _DEFAULT_DATA_DIR ) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 3 shape = (160, 160) @@ -166,7 +171,7 @@ def imagenette( def svhn( - batch_size: int, path: str = "../datasets" + batch_size: int, path: str = _DEFAULT_DATA_DIR ) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 3 shape = (32, 32) From 7264016cf55f72d5058d10c68e502ff666466810 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Tue, 14 Apr 2026 16:17:14 +0200 Subject: [PATCH 04/19] New way of creating the result directories --- scripts/symm_net/load_utils.py | 14 ++++++++++++++ scripts/symm_net/main_salnet.py | 16 +++++++++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/scripts/symm_net/load_utils.py b/scripts/symm_net/load_utils.py index a3b86ac..0e25ccb 100644 --- a/scripts/symm_net/load_utils.py +++ b/scripts/symm_net/load_utils.py @@ -85,6 +85,9 @@ def settings_loader() -> tuple[ list[str], str | None, str, + int, + str | None, + str | None, ]: parser = argparse.ArgumentParser() parser.add_argument("-f", type=str, help="Path to parameter file.") @@ -97,6 +100,14 @@ def settings_loader() -> tuple[ parser.add_argument( "--output-dir", type=str, default="../../results/symm_net", dest="output_dir" ) + parser.add_argument("--seed", type=int, default=0, help="Random seed.") + parser.add_argument( + "--run-dir", + type=str, + default=None, + dest="run_dir", + help="Exact output directory; bypasses create_run_dirs() when set.", + ) args = parser.parse_args() if args.f: @@ -137,4 +148,7 @@ def settings_loader() -> tuple[ group_tags_list, args.f, args.output_dir, + args.seed, + args.run_dir, + args.s, ) diff --git a/scripts/symm_net/main_salnet.py b/scripts/symm_net/main_salnet.py index 59ef170..f8bacba 100644 --- a/scripts/symm_net/main_salnet.py +++ b/scripts/symm_net/main_salnet.py @@ -78,8 +78,15 @@ group_tags, param_file, output_dir, + seed, + run_dir_override, + section, ) = settings_loader() +# set random seeds for reproducibility +torch.manual_seed(seed) +np.random.seed(seed) + # some general checks: assert not (params["use_backprop"] and sal_params["use_sal"]) assert not (params["use_backprop"] and rdd_params["use_rdd"]) @@ -119,7 +126,12 @@ def append_metric(metrics_dict: dict[str, Any], key: str, value: float) -> None: # Data tracking setup # --------------------------- -run_dir = create_run_dirs(base_dir=output_dir, tags=tags) +if run_dir_override is not None: + run_dir = run_dir_override + os.makedirs(os.path.join(run_dir, "figs", "weights"), exist_ok=True) + os.makedirs(os.path.join(run_dir, "checkpoints"), exist_ok=True) +else: + run_dir = create_run_dirs(base_dir=output_dir, tags=tags) # copy param file to run root if param_file is not None: @@ -131,6 +143,8 @@ def append_metric(metrics_dict: dict[str, Any], key: str, value: float) -> None: "sal_params": sal_params, "rdd_params": rdd_params, "dataset": dataset.__name__, + "algo": section, + "seed": seed, "scalars": {}, } From 960c3fc987d1f41461886bb5d13403bf9e9f79ee Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Tue, 14 Apr 2026 16:17:55 +0200 Subject: [PATCH 05/19] Add sweep launch script in python for paper runs --- scripts/symm_net/sweep.py | 112 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 scripts/symm_net/sweep.py diff --git a/scripts/symm_net/sweep.py b/scripts/symm_net/sweep.py new file mode 100644 index 0000000..8fe0fee --- /dev/null +++ b/scripts/symm_net/sweep.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Sweep launcher: runs dataset × algo × seed combinations for the paper. + +Full paper run (105 runs): + python sweep.py + +Subset examples: + python sweep.py --datasets cifar10 --algos bp sal --n-seeds 2 + python sweep.py --algos sal rdd --sweep-name sal_rdd_only +""" + +import argparse +import subprocess +import sys +from pathlib import Path + +ALL_DATASETS = ["cifar10", "fmnist", "svhn"] +ALL_ALGOS = ["bp", "fa", "bp_w_fa", "akrout", "scfa", "sal", "rdd"] +PARAM_FILE = "fast_exp.yaml" # "exp_settings.yaml" + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + p.add_argument( + "--datasets", + nargs="+", + default=ALL_DATASETS, + choices=ALL_DATASETS, + metavar="DS", + help="Datasets to run (default: all three).", + ) + p.add_argument( + "--algos", + nargs="+", + default=ALL_ALGOS, + choices=ALL_ALGOS, + metavar="ALGO", + help="Algorithm sections from exp_settings.yaml (default: all seven).", + ) + p.add_argument( + "--n-seeds", + type=int, + default=5, + dest="n_seeds", + help="Number of seeds, numbered 0 … N-1 (default: 5).", + ) + p.add_argument( + "--sweep-name", + default="sweep", + dest="sweep_name", + help="Subdirectory name under --base-dir (default: sweep).", + ) + p.add_argument( + "--base-dir", + default="../../results/symm_net", + dest="base_dir", + help="Root output directory (default: ../../results/symm_net).", + ) + return p.parse_args() + + +def main() -> None: + args = parse_args() + sweep_dir = Path(args.base_dir) / args.sweep_name + total = len(args.datasets) * len(args.algos) * args.n_seeds + done = 0 + + print(f"Sweep: {total} runs → {sweep_dir}") + print(f" datasets : {args.datasets}") + print(f" algos : {args.algos}") + print(f" seeds : 0 … {args.n_seeds - 1}") + print() + + for dataset in args.datasets: + for algo in args.algos: + for seed in range(args.n_seeds): + run_dir = sweep_dir / dataset / algo / f"seed_{seed}" + done += 1 + + if (run_dir / "metrics.json").exists(): + print(f"[{done}/{total}] Skip {dataset}/{algo}/seed_{seed}") + continue + + print(f"[{done}/{total}] Run {dataset}/{algo}/seed_{seed} ...") + result = subprocess.run( + [ + sys.executable, + "main_salnet.py", + "-f", + PARAM_FILE, + "-s", + algo, + "--dataset", + dataset, + "--seed", + str(seed), + "--run-dir", + str(run_dir), + ], + check=False, + ) + if result.returncode != 0: + print(f" WARNING: run failed (exit code {result.returncode})") + + print(f"\nSweep complete. Results in {sweep_dir}") + + +if __name__ == "__main__": + main() From cb5870ca46f2b7ff5967ce3b8d08919b7c41e423 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Tue, 14 Apr 2026 16:19:25 +0200 Subject: [PATCH 06/19] Copy over plotting script and remove neptune It's adapted to the new local sim data storage --- scripts/symm_net/plots.ipynb | 594 +++++++++++++++++++++++++++++++++++ 1 file changed, 594 insertions(+) create mode 100644 scripts/symm_net/plots.ipynb diff --git a/scripts/symm_net/plots.ipynb b/scripts/symm_net/plots.ipynb new file mode 100644 index 0000000..81b0ee3 --- /dev/null +++ b/scripts/symm_net/plots.ipynb @@ -0,0 +1,594 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Plots for symmnet" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import os\n", + "import pickle # noqa: F401 — kept for legacy compatibility if needed\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "import matplotlib as mpl\n", + "import matplotlib.pyplot as plt\n", + "import matplotlib.gridspec as gridspec" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "SAVEFIG = True\n", + "\n", + "# Point this at the sweep directory produced by sweep.py\n", + "# SWEEP_DIR = Path(\"../../results/symm_net/sweep\")\n", + "SWEEP_DIR = Path(\"/tmp/symmnet_test/sweep/\")\n", + "\n", + "FIG_DIR = Path(\"../../figs/symm_net\")\n", + "FIG_DIR.mkdir(parents=True, exist_ok=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "# define the style etc.\n", + "mpl.style.use(\"../../mystyle.mpl\")\n", + "plt.style.use(\"tableau-colorblind10\")\n", + "plt.rcParams[\"text.usetex\"] = True\n", + "mpl.rcParams[\"text.latex.preamble\"] = r\"\\usepackage{amssymb}\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "def load_runs(\n", + " sweep_dir: Path, dataset: str, algo: str, metric_key: str\n", + ") -> pd.DataFrame:\n", + " \"\"\"Load a metric time series for all seeds of one dataset/algo combination.\n", + "\n", + " Args:\n", + " sweep_dir: Root sweep directory (contains dataset sub-dirs).\n", + " dataset: Dataset name (e.g. \"cifar10\").\n", + " algo: Algorithm section name (e.g. \"sal\").\n", + " metric_key: Key inside metrics[\"scalars\"] (e.g. \"accuracy/test\").\n", + "\n", + " Returns:\n", + " DataFrame with one column per seed (named seed_0, seed_1, …) and one\n", + " row per epoch.\n", + " \"\"\"\n", + " df = pd.DataFrame()\n", + " runs_dir = sweep_dir / dataset / algo\n", + " for seed_dir in sorted(runs_dir.glob(\"seed_*\")):\n", + " path = seed_dir / \"metrics.json\"\n", + " if not path.exists():\n", + " continue\n", + " with open(path) as f:\n", + " data = json.load(f)\n", + " df[seed_dir.name] = data[\"scalars\"].get(metric_key, [])\n", + " return df" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "def add_stats(df: pd.DataFrame) -> None:\n", + " \"\"\"Add mean and std columns across all seed columns (in-place).\"\"\"\n", + " seed_cols = [c for c in df.columns if c.startswith(\"seed_\")]\n", + " df[\"mean\"] = df[seed_cols].mean(axis=1)\n", + " df[\"std\"] = df[seed_cols].std(axis=1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "def calc_final(df: pd.DataFrame, n_last: int = 5) -> tuple[float, float]:\n", + " seed_cols = [c for c in df.columns if c.startswith(\"seed_\")]\n", + " av_last = df[seed_cols].tail(n_last).mean()\n", + " return av_last.mean(), av_last.std()\n", + "\n", + "\n", + "def calc_final_all(\n", + " dfs: dict[str, pd.DataFrame], n_last: int = 5\n", + ") -> tuple[pd.DataFrame, pd.DataFrame]:\n", + " means = pd.DataFrame()\n", + " stds = pd.DataFrame()\n", + " for key, df in dfs.items():\n", + " mean, std = calc_final(df, n_last=n_last)\n", + " means[key] = [mean]\n", + " stds[key] = [std]\n", + " return means, stds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_epochs(fig, ax, df, label):\n", + " ax.plot(df[\"mean\"], label=label)\n", + " ax.fill_between(\n", + " list(range(len(df[\"mean\"]))),\n", + " df[\"mean\"] - df[\"std\"],\n", + " df[\"mean\"] + df[\"std\"],\n", + " alpha=0.3,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_bars(fig, ax, df_means, df_stds):\n", + "\n", + " # means, stds = calc_final_all(dfs)\n", + " algos = [\"BP\", \"SAL\", \"FA\", \"KP\", \"RDD\"]\n", + " means = [df_means[i] for i in algos]\n", + " stds = [df_stds[i] for i in algos]\n", + " num = len(means)\n", + " x = np.arange(num)\n", + " width = 0.8\n", + "\n", + " tab_colors = [f\"C{i}\" for i in range(7)]\n", + " fontcolors = [\"white\", \"black\", \"black\", \"white\", \"black\"]\n", + "\n", + " bars = ax.bar(x, means, width, yerr=stds, capsize=4, color=tab_colors[:num])\n", + " ax.set_xticks(x)\n", + " ax.set_xticklabels(\"\")\n", + "\n", + " TEXT_LIMIT = 75.0\n", + "\n", + " for i, (bar, fc) in enumerate(zip(bars, fontcolors)):\n", + " height = bar.get_height()\n", + " mean = means[i]\n", + " std = stds[i]\n", + " label = f\"{mean:.1f}±{std:.1f}\"\n", + " ax.text(\n", + " bar.get_x() + bar.get_width() / 2, # x-coordinate: center of the bar\n", + " (\n", + " height - std - 1.0 if height > TEXT_LIMIT else height + std + 1.0\n", + " ), # y-coordinate: slightly above the bar top, adjust as needed\n", + " label,\n", + " ha=\"center\", # horizontally center the label\n", + " va=(\n", + " \"top\" if height > TEXT_LIMIT else \"bottom\"\n", + " ), # vertically align the label to the bottom\n", + " rotation=90, # rotate the label by 45 degrees\n", + " fontsize=10,\n", + " color=fc,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "# Maps YAML section names to display labels used in the plots\n", + "ALGO_LABELS: dict[str, str] = {\n", + " \"bp\": \"BP\",\n", + " \"fa\": \"FA\",\n", + " \"bp_w_fa\": \"BP + FA\",\n", + " \"akrout\": \"KP\",\n", + " \"scfa\": \"SC FA\",\n", + " \"sal\": \"SAL\",\n", + " \"rdd\": \"RDD\",\n", + "}\n", + "\n", + "# Algorithms shown in the upper epoch-curves panel (CIFAR-10 detail)\n", + "cifar10_algos = [\"bp\", \"sal\", \"fa\", \"akrout\", \"rdd\"]\n", + "cifar10_labels = [ALGO_LABELS[a] for a in cifar10_algos]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "# Create the main figure\n", + "fig = plt.figure(figsize=(18 / 2.54, 11 / 2.54))\n", + "\n", + "# Set up main GridSpec with 2 rows (for 2 subfigures)\n", + "gs = gridspec.GridSpec(\n", + " 2,\n", + " 1,\n", + " height_ratios=[1, 1.3],\n", + " figure=fig,\n", + " top=0.85,\n", + " bottom=0.03,\n", + " right=0.95,\n", + " hspace=0.5,\n", + ")\n", + "\n", + "# Upper subfigure: 1 row, 4 columns\n", + "gs_upper = gridspec.GridSpecFromSubplotSpec(\n", + " 1, 5, subplot_spec=gs[0], width_ratios=[1.0, 0.3, 1, 1, 1], wspace=0.1\n", + ")\n", + "# axes_upper[0] bleibt wie gehabt\n", + "ax0 = fig.add_subplot(gs_upper[0, 0])\n", + "\n", + "# axes_upper[1], [2], [3] bekommen gemeinsame y-Achse\n", + "ax1 = fig.add_subplot(gs_upper[0, 2])\n", + "ax2 = fig.add_subplot(gs_upper[0, 3], sharey=ax1)\n", + "ax3 = fig.add_subplot(gs_upper[0, 4], sharey=ax1)\n", + "axes_upper = [ax0, ax1, ax2, ax3]\n", + "\n", + "# Lower subfigure: 1 row, 3 columns\n", + "gs_lower = gridspec.GridSpecFromSubplotSpec(1, 3, subplot_spec=gs[1], wspace=0.07)\n", + "ax1 = fig.add_subplot(gs_lower[0, 0])\n", + "ax2 = fig.add_subplot(gs_lower[0, 1], sharey=ax1)\n", + "ax3 = fig.add_subplot(gs_lower[0, 2], sharey=ax1)\n", + "axes_lower = [ax1, ax2, ax3]\n", + "for ax in axes_lower:\n", + " ax.spines[\"right\"].set_visible(False)\n", + " ax.spines[\"top\"].set_visible(False)\n", + "for ax in axes_lower[1:]:\n", + " ax.yaxis.set_visible(False)\n", + " ax.spines[\"left\"].set_visible(False)" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "# test accuracy:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "dfs = {\n", + " algo: load_runs(SWEEP_DIR, \"cifar10\", algo, \"accuracy/test\")\n", + " for algo in cifar10_algos\n", + "}\n", + "for df in dfs.values():\n", + " add_stats(df)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "for df, label in zip(dfs.values(), cifar10_labels):\n", + " plot_epochs(fig, axes_upper[0], df, label)\n", + "\n", + "axes_upper[0].set_xlabel(\"epochs\")\n", + "axes_upper[0].set_ylabel(\"validation accuracy [\\%]\")\n", + "axes_upper[0].set_title(\"CIFAR-10\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "fig.legend(\n", + " *axes_upper[0].get_legend_handles_labels(),\n", + " loc=\"upper center\",\n", + " bbox_to_anchor=(0.5, 0.98), # leicht unterhalb des oberen Rands\n", + " ncol=5,\n", + " frameon=True\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "15", + "metadata": {}, + "source": [ + "# angles:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [ + "dfs = {\n", + " algo: load_runs(SWEEP_DIR, \"cifar10\", algo, \"symm/angle/0\")\n", + " for algo in cifar10_algos\n", + "}\n", + "for df in dfs.values():\n", + " add_stats(df)\n", + "\n", + "for df, label in zip(dfs.values(), cifar10_labels):\n", + " plot_epochs(fig, axes_upper[1], df, label)\n", + "\n", + "axes_upper[1].set_xlabel(\"epochs\")\n", + "axes_upper[1].set_ylabel(r\"$\\measuredangle (\\mathbf W^T, \\mathbf B)\\ [\\mathrm{deg}]$\")\n", + "axes_upper[1].set_yticks([0, 45, 90])\n", + "axes_upper[1].set_title(\"FC 1\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "dfs = {\n", + " algo: load_runs(SWEEP_DIR, \"cifar10\", algo, \"symm/angle/1\")\n", + " for algo in cifar10_algos\n", + "}\n", + "for df in dfs.values():\n", + " add_stats(df)\n", + "\n", + "for df, label in zip(dfs.values(), cifar10_labels):\n", + " plot_epochs(fig, axes_upper[2], df, label)\n", + "\n", + "axes_upper[2].set_xlabel(\"epochs\")\n", + "axes_upper[2].set_title(\"FC 2\")\n", + "axes_upper[2].tick_params(axis=\"y\", labelleft=False)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "dfs = {\n", + " algo: load_runs(SWEEP_DIR, \"cifar10\", algo, \"symm/angle/2\")\n", + " for algo in cifar10_algos\n", + "}\n", + "for df in dfs.values():\n", + " add_stats(df)\n", + "\n", + "for df, label in zip(dfs.values(), cifar10_labels):\n", + " plot_epochs(fig, axes_upper[3], df, label)\n", + "\n", + "axes_upper[3].set_xlabel(\"epochs\")\n", + "axes_upper[3].set_title(\"FC 3\")\n", + "axes_upper[3].tick_params(axis=\"y\", labelleft=False)" + ] + }, + { + "cell_type": "markdown", + "id": "19", + "metadata": {}, + "source": [ + "# different data sets" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "## table with extra values in SM:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "datasets_local = [\"cifar10\", \"fmnist\", \"svhn\"]\n", + "datasets_labels = [\"CIFAR-10\", \"FMNIST\", \"SVHN\"]\n", + "# All 7 algorithms for the summary table / bar plots\n", + "all_algos = [\"bp\", \"fa\", \"bp_w_fa\", \"akrout\", \"scfa\", \"sal\", \"rdd\"]\n", + "all_labels = [ALGO_LABELS[a] for a in all_algos]\n", + "\n", + "means_list = []\n", + "stds_list = []\n", + "for dataset in datasets_local:\n", + " dfs = {\n", + " algo: load_runs(SWEEP_DIR, dataset, algo, \"accuracy/test\") for algo in all_algos\n", + " }\n", + " mean, std = calc_final_all(dfs)\n", + " mean.columns = all_labels\n", + " std.columns = all_labels\n", + " means_list.append(mean)\n", + " stds_list.append(std)\n", + "\n", + "means = pd.concat(means_list, ignore_index=True)\n", + "stds = pd.concat(stds_list, ignore_index=True)\n", + "means.index = datasets_labels\n", + "stds.index = datasets_labels" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "22", + "metadata": {}, + "outputs": [], + "source": [ + "latex_df = means.copy()\n", + "for col in means.columns:\n", + " latex_df[col] = (\n", + " means[col].round(1).astype(str) + \"$\\\\pm$\" + stds[col].round(1).astype(str)\n", + " )\n", + "\n", + "latex_df" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "with open(\"test_err_table.tex\", \"w\") as f:\n", + " f.write(latex_df.to_latex(column_format=\"l\" + \"c\" * len(latex_df.columns)))" + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": {}, + "source": [ + "## CIFAR 10" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": {}, + "outputs": [], + "source": [ + "plot_bars(fig, axes_lower[0], means.loc[\"CIFAR-10\"], stds.loc[\"CIFAR-10\"])\n", + "\n", + "bp_fa = means.loc[\"CIFAR-10\", \"BP + FA\"]\n", + "axes_lower[0].axhline(\n", + " bp_fa, color=\"green\", linestyle=\"--\", label=\"theo. upper limit\", zorder=-1\n", + ")\n", + "\n", + "axes_lower[0].set_ylim(60, 100)\n", + "axes_lower[0].set_title(\"CIFAR-10\", pad=0)\n", + "axes_lower[0].set_ylabel(\"test accuracy [\\%]\")\n", + "axes_lower[0].legend()" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, + "source": [ + "## FMNIST" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27", + "metadata": {}, + "outputs": [], + "source": [ + "plot_bars(fig, axes_lower[1], means.loc[\"FMNIST\"], stds.loc[\"FMNIST\"])\n", + "\n", + "bp_fa = means.loc[\"FMNIST\", \"BP + FA\"]\n", + "axes_lower[1].axhline(bp_fa, color=\"green\", linestyle=\"--\", zorder=-1)\n", + "\n", + "axes_lower[1].set_title(\"Fashion-MNIST\", pad=0)\n", + "axes_lower[1].tick_params(axis=\"y\", labelleft=False)" + ] + }, + { + "cell_type": "markdown", + "id": "28", + "metadata": {}, + "source": [ + "## SVHN" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29", + "metadata": {}, + "outputs": [], + "source": [ + "plot_bars(fig, axes_lower[2], means.loc[\"SVHN\"], stds.loc[\"SVHN\"])\n", + "\n", + "bp_fa = means.loc[\"SVHN\", \"BP + FA\"]\n", + "axes_lower[2].axhline(bp_fa, color=\"green\", linestyle=\"--\", zorder=-1)\n", + "\n", + "axes_lower[2].set_title(\"SVHN\", pad=0)\n", + "axes_lower[2].tick_params(axis=\"y\", labelleft=False)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "30", + "metadata": {}, + "outputs": [], + "source": [ + "display(fig)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "31", + "metadata": {}, + "outputs": [], + "source": [ + "if SAVEFIG:\n", + " fig.savefig(\"../figs/salnet.png\", dpi=300)\n", + " fig.savefig(\"../figs/salnet.pdf\")\n", + " fig.savefig(\"../figs/salnet.svg\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "sal", + "language": "python", + "name": "sal" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.10" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From b1d3ac1944db0273571f542e34380d16c2d57bb8 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Tue, 14 Apr 2026 16:57:29 +0200 Subject: [PATCH 07/19] Nicefy sweep.py (a bit) --- scripts/symm_net/sweep.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/scripts/symm_net/sweep.py b/scripts/symm_net/sweep.py index 8fe0fee..4c220f1 100644 --- a/scripts/symm_net/sweep.py +++ b/scripts/symm_net/sweep.py @@ -69,9 +69,9 @@ def main() -> None: done = 0 print(f"Sweep: {total} runs → {sweep_dir}") - print(f" datasets : {args.datasets}") - print(f" algos : {args.algos}") - print(f" seeds : 0 … {args.n_seeds - 1}") + print(f" datasets: {args.datasets}") + print(f" algos: {args.algos}") + print(f" seeds: 0 … {args.n_seeds - 1}") print() for dataset in args.datasets: From b07f7b97e039cf72bf872f7900b6579dd13d1253 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Wed, 15 Apr 2026 11:15:58 +0200 Subject: [PATCH 08/19] Add cuncurrent runners to sweep.py --- scripts/symm_net/sweep.py | 114 ++++++++++++++++++++++++++------------ 1 file changed, 78 insertions(+), 36 deletions(-) diff --git a/scripts/symm_net/sweep.py b/scripts/symm_net/sweep.py index 4c220f1..9ea1360 100644 --- a/scripts/symm_net/sweep.py +++ b/scripts/symm_net/sweep.py @@ -1,22 +1,30 @@ #!/usr/bin/env python3 -"""Sweep launcher: runs dataset × algo × seed combinations for the paper. +"""Sweep launcher for convenient reproduction of data for SymmNet paper figure -Full paper run (105 runs): +Runs dataset × algo × seed combinations. + +Sequential run (default): python sweep.py -Subset examples: +Parallel run with 4 workers: + python sweep.py --n-workers 4 + +Subset example: python sweep.py --datasets cifar10 --algos bp sal --n-seeds 2 - python sweep.py --algos sal rdd --sweep-name sal_rdd_only + +Re-running is safe: finished runs (metrics.json present) are skipped. """ import argparse import subprocess import sys +from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path ALL_DATASETS = ["cifar10", "fmnist", "svhn"] ALL_ALGOS = ["bp", "fa", "bp_w_fa", "akrout", "scfa", "sal", "rdd"] -PARAM_FILE = "fast_exp.yaml" # "exp_settings.yaml" +# PARAM_FILE = "exp_settings.yaml" +PARAM_FILE = "fast_exp.yaml" def parse_args() -> argparse.Namespace: @@ -47,6 +55,13 @@ def parse_args() -> argparse.Namespace: dest="n_seeds", help="Number of seeds, numbered 0 … N-1 (default: 5).", ) + p.add_argument( + "--n-workers", + type=int, + default=1, + dest="n_workers", + help="Number of parallel workers (default: 1 = sequential).", + ) p.add_argument( "--sweep-name", default="sweep", @@ -62,48 +77,75 @@ def parse_args() -> argparse.Namespace: return p.parse_args() +def _run_one( + run_dir: Path, + dataset: str, + algo: str, + seed: int, + param_file: str, +) -> tuple[str, int | None]: + """Run a single training job. + + Returns: + A (label, returncode) tuple. returncode is None if the run was skipped. + """ + label = f"{dataset}/{algo}/seed_{seed}" + if (run_dir / "metrics.json").exists(): + return label, None # already done + result = subprocess.run( + [ + sys.executable, + "main_salnet.py", + "-f", + param_file, + "-s", + algo, + "--dataset", + dataset, + "--seed", + str(seed), + "--run-dir", + str(run_dir), + ], + check=False, + ) + return label, result.returncode + + def main() -> None: args = parse_args() sweep_dir = Path(args.base_dir) / args.sweep_name total = len(args.datasets) * len(args.algos) * args.n_seeds - done = 0 print(f"Sweep: {total} runs → {sweep_dir}") print(f" datasets: {args.datasets}") print(f" algos: {args.algos}") print(f" seeds: 0 … {args.n_seeds - 1}") + print(f" workers: {args.n_workers}") print() - for dataset in args.datasets: - for algo in args.algos: - for seed in range(args.n_seeds): - run_dir = sweep_dir / dataset / algo / f"seed_{seed}" - done += 1 - - if (run_dir / "metrics.json").exists(): - print(f"[{done}/{total}] Skip {dataset}/{algo}/seed_{seed}") - continue - - print(f"[{done}/{total}] Run {dataset}/{algo}/seed_{seed} ...") - result = subprocess.run( - [ - sys.executable, - "main_salnet.py", - "-f", - PARAM_FILE, - "-s", - algo, - "--dataset", - dataset, - "--seed", - str(seed), - "--run-dir", - str(run_dir), - ], - check=False, - ) - if result.returncode != 0: - print(f" WARNING: run failed (exit code {result.returncode})") + runs = [ + (sweep_dir / dataset / algo / f"seed_{seed}", dataset, algo, seed) + for dataset in args.datasets + for algo in args.algos + for seed in range(args.n_seeds) + ] + + done = 0 + with ThreadPoolExecutor(max_workers=args.n_workers) as pool: + futures = { + pool.submit(_run_one, run_dir, dataset, algo, seed, PARAM_FILE): i + for i, (run_dir, dataset, algo, seed) in enumerate(runs) + } + for future in as_completed(futures): + label, code = future.result() + done += 1 + if code is None: + print(f"[{done}/{total}] Skip {label}") + elif code != 0: + print(f"[{done}/{total}] FAIL {label} (exit {code})") + else: + print(f"[{done}/{total}] Done {label}") print(f"\nSweep complete. Results in {sweep_dir}") From ba8b35aff330c7f55b444932f433f0b717779202 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Wed, 15 Apr 2026 11:17:04 +0200 Subject: [PATCH 09/19] Add a bash sweep launcher file and a fast exp yaml ... the fast yaml should later be removed again --- scripts/microcircuits/fast_example.yaml | 101 ++++++++++++++++++++++++ scripts/symm_net/sweep.sh | 75 ++++++++++++++++++ 2 files changed, 176 insertions(+) create mode 100644 scripts/microcircuits/fast_example.yaml create mode 100644 scripts/symm_net/sweep.sh diff --git a/scripts/microcircuits/fast_example.yaml b/scripts/microcircuits/fast_example.yaml new file mode 100644 index 0000000..5206cd9 --- /dev/null +++ b/scripts/microcircuits/fast_example.yaml @@ -0,0 +1,101 @@ +exp_type: "test workflow" +network_properties: + dims: + - 1 + - 1 + - 1 + lambda_api: 0.02 + lambda_nudge: 0.6 + learning_lag: 2000 + lr: + - down: 0.001 + up: 0.4 + - up: 0.001 + n_last_spks: 1 + size_moving_average: 2000 + stdp_a_anticausal: 1.0 + stdp_a_causal: -1.0 + stdp_tau_anticausal: 10.0 + stdp_tau_causal: 10.0 + t_ref: 10 + tau_syn: 10 +student_initial_parameters: + bias: + - -1.0 + - -1.0 + share_weights: false + random_weights_init_limits: + up: + - -3.0 + - 3.0 + down: + - -3.0 + - 3.0 +student_simulation_settings: + len_epoch: 5 + len_symmetrization: 2 + len_validation: 4 + num_epochs: 2 + poisson_seed: 649854 + weights_init_seed: 355445 + record_all_spks: false + recorded_quantities: + - - u_in + - - u_pyr + - v_api + - w_up + - w_down + - - u_pyr + - v_nudge + - u_tgt + - w_up + - validation + - symmetrization + recorded_sample_length: 20000 + set_sps: true + shuffle_training: true + shuffle_validation: false + t_pattern: 2000 + training_seed: 915344 + update_down: true +teacher_initial_parameters: + bias: + - -1.0 + - -1.0 + weights: + - w_up: 2.0 + - w_up: 2.0 +teacher_simulation_settings: + len_epoch: 8 + len_validation: 6 + num_epochs: 1 + poisson_seed: 423433 + recorded_quantities: + - - u_in + - - u_pyr + - v_bas + - v_api + - u_inn + - w_up + - w_down + - - u_pyr + - v_bas + - v_nudge + - u_tgt + - w_up + shuffle_training: true + shuffle_validation: false + t_pattern: 2000 + training_seed: 890423 +u_input: + training: + - -2.0 + - -0.8 + - 0.8 + - 2.0 + - 3.0 + validation: + - -1.8 + - -0.5 + - 1.8 + - 2.7 diff --git a/scripts/symm_net/sweep.sh b/scripts/symm_net/sweep.sh new file mode 100644 index 0000000..0b5e8c2 --- /dev/null +++ b/scripts/symm_net/sweep.sh @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +# sweep.sh — run all dataset × algo × seed combinations for the paper. +# +# Usage (sequential local run): +# bash sweep.sh +# +# For SLURM: replace the `python` call in the marked block below with an +# sbatch invocation. Minimal example: +# +# sbatch \ +# --partition=gpu --gres=gpu:1 --time=4:00:00 --mem=16G \ +# --job-name="${ALGO}_${DATASET}_s${SEED}" \ +# --wrap="cd $(pwd) && conda activate myenv && python main_salnet.py \ +# -f ${PARAM_FILE} -s ${ALGO} --dataset ${DATASET} \ +# --seed ${SEED} --run-dir ${RUN_DIR}" +# +# Re-running is safe: finished runs (metrics.json present) are skipped. + +set -euo pipefail + +# ── Configuration ──────────────────────────────────────────────────────────── +DATASETS=("cifar10" "fmnist" "svhn") +ALGOS=("bp" "fa" "bp_w_fa" "akrout" "scfa" "sal" "rdd") +N_SEEDS=5 +SWEEP_NAME="sweep" +BASE_DIR="../../results/symm_net" +PARAM_FILE="exp_settings.yaml" +# ───────────────────────────────────────────────────────────────────────────── + +DATASETS=("fmnist") +ALGOS=("bp") +N_SEEDS=2 +SWEEP_NAME="sh_sweep_1" +BASE_DIR="/tmp/symmnet_test" +PARAM_FILE="fast_exp.yaml" + +SWEEP_DIR="${BASE_DIR}/${SWEEP_NAME}" +TOTAL=$(( ${#DATASETS[@]} * ${#ALGOS[@]} * N_SEEDS )) +DONE=0 + +echo "Sweep: ${TOTAL} runs → ${SWEEP_DIR}" +echo " datasets: ${DATASETS[*]}" +echo " algos: ${ALGOS[*]}" +echo " seeds: 0 … $(( N_SEEDS - 1 ))" +echo "" + +for DATASET in "${DATASETS[@]}"; do + for ALGO in "${ALGOS[@]}"; do + for (( SEED=0; SEED Date: Wed, 15 Apr 2026 12:23:43 +0200 Subject: [PATCH 10/19] BUGFIX: add the corrrect fast exp file ... and remove the other --- scripts/microcircuits/fast_example.yaml | 101 ------------------------ scripts/symm_net/fast_exp.yaml | 94 ++++++++++++++++++++++ 2 files changed, 94 insertions(+), 101 deletions(-) delete mode 100644 scripts/microcircuits/fast_example.yaml create mode 100644 scripts/symm_net/fast_exp.yaml diff --git a/scripts/microcircuits/fast_example.yaml b/scripts/microcircuits/fast_example.yaml deleted file mode 100644 index 5206cd9..0000000 --- a/scripts/microcircuits/fast_example.yaml +++ /dev/null @@ -1,101 +0,0 @@ -exp_type: "test workflow" -network_properties: - dims: - - 1 - - 1 - - 1 - lambda_api: 0.02 - lambda_nudge: 0.6 - learning_lag: 2000 - lr: - - down: 0.001 - up: 0.4 - - up: 0.001 - n_last_spks: 1 - size_moving_average: 2000 - stdp_a_anticausal: 1.0 - stdp_a_causal: -1.0 - stdp_tau_anticausal: 10.0 - stdp_tau_causal: 10.0 - t_ref: 10 - tau_syn: 10 -student_initial_parameters: - bias: - - -1.0 - - -1.0 - share_weights: false - random_weights_init_limits: - up: - - -3.0 - - 3.0 - down: - - -3.0 - - 3.0 -student_simulation_settings: - len_epoch: 5 - len_symmetrization: 2 - len_validation: 4 - num_epochs: 2 - poisson_seed: 649854 - weights_init_seed: 355445 - record_all_spks: false - recorded_quantities: - - - u_in - - - u_pyr - - v_api - - w_up - - w_down - - - u_pyr - - v_nudge - - u_tgt - - w_up - - validation - - symmetrization - recorded_sample_length: 20000 - set_sps: true - shuffle_training: true - shuffle_validation: false - t_pattern: 2000 - training_seed: 915344 - update_down: true -teacher_initial_parameters: - bias: - - -1.0 - - -1.0 - weights: - - w_up: 2.0 - - w_up: 2.0 -teacher_simulation_settings: - len_epoch: 8 - len_validation: 6 - num_epochs: 1 - poisson_seed: 423433 - recorded_quantities: - - - u_in - - - u_pyr - - v_bas - - v_api - - u_inn - - w_up - - w_down - - - u_pyr - - v_bas - - v_nudge - - u_tgt - - w_up - shuffle_training: true - shuffle_validation: false - t_pattern: 2000 - training_seed: 890423 -u_input: - training: - - -2.0 - - -0.8 - - 0.8 - - 2.0 - - 3.0 - validation: - - -1.8 - - -0.5 - - 1.8 - - 2.7 diff --git a/scripts/symm_net/fast_exp.yaml b/scripts/symm_net/fast_exp.yaml new file mode 100644 index 0000000..c03ae52 --- /dev/null +++ b/scripts/symm_net/fast_exp.yaml @@ -0,0 +1,94 @@ +bp: + params: + n_epochs: 1 + batch_size: 128 + lr: 0.01 + momentum: 0.9 + weight_decay: 0.0 + use_backprop: True + use_kp: False + use_fa_conv_layers: False + use_scfa: False + +fa: + params: + n_epochs: 1 + batch_size: 128 + lr: 0.01 + momentum: 0.9 + weight_decay: 0.0 + use_backprop: False + use_kp: False + use_fa_conv_layers: False + use_scfa: False + +bp_w_fa: + params: + n_epochs: 1 + batch_size: 128 + lr: 0.01 + momentum: 0.9 + weight_decay: 0.0 + use_backprop: True + use_kp: False + use_fa_conv_layers: True + use_scfa: False + +akrout: + params: + n_epochs: 1 + batch_size: 128 + lr: 0.01 + momentum: 0.9 + weight_decay: 0.001 + use_backprop: False + use_kp: True + use_fa_conv_layers: False + use_scfa: False + +scfa: + params: + n_epochs: 1 + batch_size: 128 + lr: 0.01 + momentum: 0.9 + weight_decay: 0.0 + use_backprop: False + use_kp: False + use_fa_conv_layers: False + use_scfa: True + +sal: + params: + n_epochs: 1 + batch_size: 128 + lr: 0.01 + momentum: 0.9 + weight_decay: 0.0 + use_backprop: False + use_kp: False + use_fa_conv_layers: False + use_scfa: False + sal_params: + n_iterations: 5 + use_sal: True + t_ref: 10 + len_epoch: 200 + sal_lr: 0.04 + batch_size: 32 + +rdd: + params: + n_epochs: 1 + batch_size: 128 + lr: 0.01 + momentum: 0.9 + weight_decay: 0.0 + use_backprop: False + use_kp: False + use_fa_conv_layers: False + use_scfa: False + rdd_params: + rdd_time: 90 + use_rdd: True + every_epoch: True From 6c818f97170e018df669f9b6dd7a6da8d3e61b50 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Wed, 15 Apr 2026 17:08:06 +0200 Subject: [PATCH 11/19] Add slurm compatible alternative for job spawning --- scripts/symm_net/sweep_creator.py | 112 ++++++++++++++++++++++++++++++ scripts/symm_net/sweep_starter.sh | 16 +++++ 2 files changed, 128 insertions(+) create mode 100644 scripts/symm_net/sweep_creator.py create mode 100644 scripts/symm_net/sweep_starter.sh diff --git a/scripts/symm_net/sweep_creator.py b/scripts/symm_net/sweep_creator.py new file mode 100644 index 0000000..ffb45fb --- /dev/null +++ b/scripts/symm_net/sweep_creator.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Sweep launcher for convenient reproduction of data for SymmNet paper figure + +Runs dataset × algo × seed combinations. + +Sequential run (default): + python sweep.py + +Parallel run with 4 workers: + python sweep.py --n-workers 4 + +Subset example: + python sweep.py --datasets cifar10 --algos bp sal --n-seeds 2 + +Re-running is safe: finished runs (metrics.json present) are skipped. +""" + +import argparse +from pathlib import Path + +ALL_DATASETS = ["cifar10", "fmnist", "svhn"] +ALL_ALGOS = ["bp", "fa", "bp_w_fa", "akrout", "scfa", "sal", "rdd"] +# PARAM_FILE = "exp_settings.yaml" +PARAM_FILE = "fast_exp.yaml" +PY_FILE = "main_salnet.py" +JOBS_FILE = "jobs.sh" + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + p.add_argument( + "--datasets", + nargs="+", + default=ALL_DATASETS, + choices=ALL_DATASETS, + metavar="DS", + help="Datasets to run (default: all three).", + ) + p.add_argument( + "--algos", + nargs="+", + default=ALL_ALGOS, + choices=ALL_ALGOS, + metavar="ALGO", + help="Algorithm sections from exp_settings.yaml (default: all seven).", + ) + p.add_argument( + "--n-seeds", + type=int, + default=5, + dest="n_seeds", + help="Number of seeds, numbered 0 … N-1 (default: 5).", + ) + p.add_argument( + "--sweep-name", + default="sweep", + dest="sweep_name", + help="Subdirectory name under --base-dir (default: sweep).", + ) + p.add_argument( + "--base-dir", + default="../../results/symm_net", + dest="base_dir", + help="Root output directory (default: ../../results/symm_net).", + ) + return p.parse_args() + + +def main() -> None: + args = parse_args() + sweep_dir = Path(args.base_dir) / args.sweep_name + total = len(args.datasets) * len(args.algos) * args.n_seeds + done = 0 + + print(f"Sweep: {total} runs → {sweep_dir}") + print(f" datasets: {args.datasets}") + print(f" algos: {args.algos}") + print(f" seeds: 0 … {args.n_seeds - 1}") + print() + + if Path(JOBS_FILE).exists(): + print(f"{JOBS_FILE} already exists. Stop here!") + return + + proc_calls = [] + for dataset in args.datasets: + for algo in args.algos: + for seed in range(args.n_seeds): + run_dir = sweep_dir / dataset / algo / f"seed_{seed}" + done += 1 + + if (run_dir / "metrics.json").exists(): + print(f"[{done}/{total}] Skip {dataset}/{algo}/seed_{seed}") + continue + + print(f"[{done}/{total}] Run {dataset}/{algo}/seed_{seed}") + proc_calls.append( + ( + f"python {PY_FILE} -f {PARAM_FILE} -s {algo}" + f" --dataset {dataset} --run-dir {run_dir}\n" + ) + ) + + with open(JOBS_FILE, "w") as f: + f.writelines(proc_calls) + + +if __name__ == "__main__": + main() diff --git a/scripts/symm_net/sweep_starter.sh b/scripts/symm_net/sweep_starter.sh new file mode 100644 index 0000000..dc5c2dd --- /dev/null +++ b/scripts/symm_net/sweep_starter.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env sh + + +set -euo pipefail + +JOBS_FILE="${1:-jobs.sh}" # Pfad als erstes Argument, default: jobs.sh + +TOTAL=$(wc -l < "${JOBS_FILE}") +echo "$TOTAL" +DONE=0 + +while IFS= read -r CMD; do + DONE=$(( DONE + 1 )) + echo "[${DONE}/${TOTAL}] ${CMD}" + eval "$CMD" +done < "${JOBS_FILE}" From 3af4e9b815fa557c6526659f27c7ed2bbe21adfc Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Thu, 16 Apr 2026 11:17:30 +0200 Subject: [PATCH 12/19] Puzzle out perfect slurm workflow --- scripts/symm_net/slurm.sh | 21 +++++++++++++++++++++ scripts/symm_net/slurm_submit.sh | 13 +++++++++++++ scripts/symm_net/sweep_creator.py | 3 +-- 3 files changed, 35 insertions(+), 2 deletions(-) create mode 100644 scripts/symm_net/slurm.sh create mode 100644 scripts/symm_net/slurm_submit.sh diff --git a/scripts/symm_net/slurm.sh b/scripts/symm_net/slurm.sh new file mode 100644 index 0000000..87a0066 --- /dev/null +++ b/scripts/symm_net/slurm.sh @@ -0,0 +1,21 @@ +#!/bin/bash +#SBATCH --job-name="symmnet sweep" +#SBATCH --time=10:00:00 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem-per-cpu=4G +#SBATCH --partition=gpu +#SBATCH --gres=gpu:rtx4090:1 +#SBATCH --account=paygo +#SBATCH --wckey=pyl_biodl + +# Your code below this line +module load Anaconda3 +eval "$(conda shell.bash hook)" + +conda activate ../../envs/ + +echo "BLA!" + +CMD=$(sed -n "${SLURM_ARRAY_TASK_ID}p" "${JOBS_FILE}") +eval "$CMD" diff --git a/scripts/symm_net/slurm_submit.sh b/scripts/symm_net/slurm_submit.sh new file mode 100644 index 0000000..b1a9234 --- /dev/null +++ b/scripts/symm_net/slurm_submit.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +# slurm_submit.sh + +# snapshot a copy of jobs.sh +TIMESTAMP=$(date +%Y%m%d_%H%M%S) +SNAPSHOT="jobs_${TIMESTAMP}.sh" + +cp jobs.sh "${SNAPSHOT}" +echo "Copy a snapshot of jobs.sh to ${SNAPSHOT}" +TOTAL=$(wc -l < "${SNAPSHOT}") + +sbatch --array=1-${TOTAL} --export="ALL,JOBS_FILE=${SNAPSHOT}" slurm.sh +echo "Submitted ${TOTAL} jobs." diff --git a/scripts/symm_net/sweep_creator.py b/scripts/symm_net/sweep_creator.py index ffb45fb..7438760 100644 --- a/scripts/symm_net/sweep_creator.py +++ b/scripts/symm_net/sweep_creator.py @@ -20,8 +20,7 @@ ALL_DATASETS = ["cifar10", "fmnist", "svhn"] ALL_ALGOS = ["bp", "fa", "bp_w_fa", "akrout", "scfa", "sal", "rdd"] -# PARAM_FILE = "exp_settings.yaml" -PARAM_FILE = "fast_exp.yaml" +PARAM_FILE = "exp_settings.yaml" PY_FILE = "main_salnet.py" JOBS_FILE = "jobs.sh" From 64364570455dd374eb20c2c088242defff12ddeb Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Thu, 16 Apr 2026 11:47:54 +0200 Subject: [PATCH 13/19] Nicify slurm sweep scripts and add dscription --- scripts/symm_net/.gitignore | 2 ++ scripts/symm_net/slurm.sh | 28 ++++++++++++------------- scripts/symm_net/slurm_submit.sh | 16 +++++++++++++- scripts/symm_net/sweep_creator.py | 35 ++++++++++++++++++++++--------- 4 files changed, 56 insertions(+), 25 deletions(-) create mode 100644 scripts/symm_net/.gitignore diff --git a/scripts/symm_net/.gitignore b/scripts/symm_net/.gitignore new file mode 100644 index 0000000..921fcb0 --- /dev/null +++ b/scripts/symm_net/.gitignore @@ -0,0 +1,2 @@ +jobs.sh +jobs_[0-9]*_[0-9]*.sh diff --git a/scripts/symm_net/slurm.sh b/scripts/symm_net/slurm.sh index 87a0066..8bb87f6 100644 --- a/scripts/symm_net/slurm.sh +++ b/scripts/symm_net/slurm.sh @@ -1,21 +1,21 @@ #!/bin/bash +# --- HPC cluster specific SLURM settings ------------------------ +# add your settings here. #SBATCH --job-name="symmnet sweep" -#SBATCH --time=10:00:00 -#SBATCH --ntasks=1 -#SBATCH --cpus-per-task=4 -#SBATCH --mem-per-cpu=4G -#SBATCH --partition=gpu -#SBATCH --gres=gpu:rtx4090:1 -#SBATCH --account=paygo -#SBATCH --wckey=pyl_biodl +#SBATCH --time= +#SBATCH ... +# ---------------------------------------------------------------- -# Your code below this line -module load Anaconda3 -eval "$(conda shell.bash hook)" +# --- HPC cluster specific setup:--------------------------------- +# e.g. load modules +# and activate python environment +# ---------------------------------------------------------------- -conda activate ../../envs/ - -echo "BLA!" +echo "Start job with id ${SLURM_ARRAY_TASK_ID}." +# read the snapshotted jobs_.sh file CMD=$(sed -n "${SLURM_ARRAY_TASK_ID}p" "${JOBS_FILE}") +echo "Run command ${CMD}" +# and execute it. eval "$CMD" +echo "Done." diff --git a/scripts/symm_net/slurm_submit.sh b/scripts/symm_net/slurm_submit.sh index b1a9234..963963a 100644 --- a/scripts/symm_net/slurm_submit.sh +++ b/scripts/symm_net/slurm_submit.sh @@ -1,5 +1,19 @@ #!/usr/bin/env bash -# slurm_submit.sh +# slurm_submit.sh — Step 2 of the SLURM sweep workflow. +# +# Prerequisites: +# jobs.sh must exist (created by sweep_creator.py). +# slurm.sh must exist and contain your #SBATCH directives. +# +# What this script does: +# 1. Snapshots jobs.sh to jobs_.sh so that late-running array +# tasks read a stable file even if jobs.sh is regenerated in the meantime. +# 2. Submits a SLURM array job (one task per line in the snapshot). +# Each task reads its command from the snapshot via $SLURM_ARRAY_TASK_ID +# and executes it (see slurm.sh). +# +# Usage: +# bash slurm_submit.sh # snapshot a copy of jobs.sh TIMESTAMP=$(date +%Y%m%d_%H%M%S) diff --git a/scripts/symm_net/sweep_creator.py b/scripts/symm_net/sweep_creator.py index 7438760..8a64b64 100644 --- a/scripts/symm_net/sweep_creator.py +++ b/scripts/symm_net/sweep_creator.py @@ -1,18 +1,33 @@ #!/usr/bin/env python3 -"""Sweep launcher for convenient reproduction of data for SymmNet paper figure +"""Generate a jobs.sh file for SLURM array job submission. -Runs dataset × algo × seed combinations. +This is step 1 of the SLURM workflow for reproducing the SymmNet paper figure. +Each line of the output file contains one `python main_salnet.py ...` command, +one per dataset × algo × seed combination. Already-completed runs (those with a +metrics.json in the expected output directory) are skipped. -Sequential run (default): - python sweep.py +Full workflow +------------- +1. Generate the job list:: -Parallel run with 4 workers: - python sweep.py --n-workers 4 + python sweep_creator.py [--datasets ...] [--algos ...] [--n-seeds N] -Subset example: - python sweep.py --datasets cifar10 --algos bp sal --n-seeds 2 + This writes ``jobs.sh`` in the current directory. -Re-running is safe: finished runs (metrics.json present) are skipped. +2. Submit as a SLURM array job (snapshots jobs.sh first for reproducibility):: + + bash slurm_submit.sh + + Internally, slurm_submit.sh calls ``sbatch --array=1-N slurm.sh``. + Each array task reads its own command from the snapshot and executes it. + +3. After all jobs finish, run ``plots.ipynb`` to reproduce the figure. + +Notes +----- +- ``jobs.sh`` must not exist before running this script (guard against + accidental overwrites); delete it manually to regenerate. +- To run locally instead of on SLURM, use ``sweep.py`` directly. """ import argparse @@ -99,7 +114,7 @@ def main() -> None: proc_calls.append( ( f"python {PY_FILE} -f {PARAM_FILE} -s {algo}" - f" --dataset {dataset} --run-dir {run_dir}\n" + f" --dataset {dataset} --seed {seed} --run-dir {run_dir}\n" ) ) From 1009e63a9d7751782bbca289b1ce31257828cab7 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Thu, 16 Apr 2026 11:48:42 +0200 Subject: [PATCH 14/19] Remove obsolete sweep scripts --- scripts/symm_net/sweep.sh | 75 ------------------------------- scripts/symm_net/sweep_starter.sh | 16 ------- 2 files changed, 91 deletions(-) delete mode 100644 scripts/symm_net/sweep.sh delete mode 100644 scripts/symm_net/sweep_starter.sh diff --git a/scripts/symm_net/sweep.sh b/scripts/symm_net/sweep.sh deleted file mode 100644 index 0b5e8c2..0000000 --- a/scripts/symm_net/sweep.sh +++ /dev/null @@ -1,75 +0,0 @@ -#!/usr/bin/env bash -# sweep.sh — run all dataset × algo × seed combinations for the paper. -# -# Usage (sequential local run): -# bash sweep.sh -# -# For SLURM: replace the `python` call in the marked block below with an -# sbatch invocation. Minimal example: -# -# sbatch \ -# --partition=gpu --gres=gpu:1 --time=4:00:00 --mem=16G \ -# --job-name="${ALGO}_${DATASET}_s${SEED}" \ -# --wrap="cd $(pwd) && conda activate myenv && python main_salnet.py \ -# -f ${PARAM_FILE} -s ${ALGO} --dataset ${DATASET} \ -# --seed ${SEED} --run-dir ${RUN_DIR}" -# -# Re-running is safe: finished runs (metrics.json present) are skipped. - -set -euo pipefail - -# ── Configuration ──────────────────────────────────────────────────────────── -DATASETS=("cifar10" "fmnist" "svhn") -ALGOS=("bp" "fa" "bp_w_fa" "akrout" "scfa" "sal" "rdd") -N_SEEDS=5 -SWEEP_NAME="sweep" -BASE_DIR="../../results/symm_net" -PARAM_FILE="exp_settings.yaml" -# ───────────────────────────────────────────────────────────────────────────── - -DATASETS=("fmnist") -ALGOS=("bp") -N_SEEDS=2 -SWEEP_NAME="sh_sweep_1" -BASE_DIR="/tmp/symmnet_test" -PARAM_FILE="fast_exp.yaml" - -SWEEP_DIR="${BASE_DIR}/${SWEEP_NAME}" -TOTAL=$(( ${#DATASETS[@]} * ${#ALGOS[@]} * N_SEEDS )) -DONE=0 - -echo "Sweep: ${TOTAL} runs → ${SWEEP_DIR}" -echo " datasets: ${DATASETS[*]}" -echo " algos: ${ALGOS[*]}" -echo " seeds: 0 … $(( N_SEEDS - 1 ))" -echo "" - -for DATASET in "${DATASETS[@]}"; do - for ALGO in "${ALGOS[@]}"; do - for (( SEED=0; SEED Date: Thu, 16 Apr 2026 15:02:03 +0200 Subject: [PATCH 15/19] Adapt salnet_symm experiments and plot script ... new local storage system --- scripts/symm_net/plot_puresymm.ipynb | 267 +++++++++++++++++++++++++ scripts/symm_net/salnet_symm.py | 13 +- scripts/symm_net/sweep.py | 2 +- scripts/symm_net/sweep_creator_symm.py | 128 ++++++++++++ scripts/symm_net/sweep_symm.py | 157 +++++++++++++++ 5 files changed, 564 insertions(+), 3 deletions(-) create mode 100644 scripts/symm_net/plot_puresymm.ipynb create mode 100644 scripts/symm_net/sweep_creator_symm.py create mode 100644 scripts/symm_net/sweep_symm.py diff --git a/scripts/symm_net/plot_puresymm.ipynb b/scripts/symm_net/plot_puresymm.ipynb new file mode 100644 index 0000000..5c8131d --- /dev/null +++ b/scripts/symm_net/plot_puresymm.ipynb @@ -0,0 +1,267 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Plots for pure symmetrization in SALNet" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "import matplotlib as mpl\n", + "import matplotlib.ticker as ticker\n", + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "SAVEFIG = True\n", + "\n", + "# Point this at the sweep directory produced by sweep_symm.py\n", + "SWEEP_DIR = Path(\"../../results/symm_net/puresymm/puresymm\")\n", + "\n", + "FIG_DIR = Path(\"../../figs/symm_net\")\n", + "FIG_DIR.mkdir(parents=True, exist_ok=True)\n", + "\n", + "# Learning rates — must match what was passed to sweep_symm.py\n", + "LRS = [0.01, 0.02, 0.04, 0.08]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "# define the style etc.\n", + "mpl.style.use(\"../../mystyle.mpl\")\n", + "plt.style.use(\"tableau-colorblind10\")\n", + "plt.rcParams[\"text.usetex\"] = True\n", + "mpl.rcParams[\"text.latex.preamble\"] = r\"\\usepackage{amssymb}\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "def load_runs(sweep_dir: Path, lr: float, metric_key: str) -> pd.DataFrame:\n", + " \"\"\"Load a metric time series for all seeds of one learning rate.\n", + "\n", + " Args:\n", + " sweep_dir: Root sweep directory (contains lr_* sub-dirs).\n", + " lr: Learning rate value.\n", + " metric_key: Key inside metrics[\"scalars\"] (e.g. \"symm/angle/0\").\n", + "\n", + " Returns:\n", + " DataFrame with one column per seed (seed_0, seed_1, …) and one row\n", + " per epoch.\n", + " \"\"\"\n", + " df = pd.DataFrame()\n", + " lr_dir = sweep_dir / f\"lr_{lr}\"\n", + " for seed_dir in sorted(lr_dir.glob(\"seed_*\")):\n", + " path = seed_dir / \"metrics.json\"\n", + " if not path.exists():\n", + " continue\n", + " with open(path) as f:\n", + " data = json.load(f)\n", + " df[seed_dir.name] = data[\"scalars\"].get(metric_key, [])\n", + " return df\n", + "\n", + "\n", + "def load_params(sweep_dir: Path, lr: float) -> dict:\n", + " \"\"\"Load params dict from the first available seed of a learning rate.\"\"\"\n", + " lr_dir = sweep_dir / f\"lr_{lr}\"\n", + " for seed_dir in sorted(lr_dir.glob(\"seed_*\")):\n", + " path = seed_dir / \"metrics.json\"\n", + " if path.exists():\n", + " with open(path) as f:\n", + " return json.load(f)[\"params\"]\n", + " return {}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "dfs_1, dfs_2, dfs_3, lr_labels = [], [], [], []\n", + "for lr in LRS:\n", + " n_seeds = len(list((SWEEP_DIR / f\"lr_{lr}\").glob(\"seed_*\")))\n", + " dfs_1.append(load_runs(SWEEP_DIR, lr, \"symm/angle/0\"))\n", + " dfs_2.append(load_runs(SWEEP_DIR, lr, \"symm/angle/1\"))\n", + " dfs_3.append(load_runs(SWEEP_DIR, lr, \"symm/angle/2\"))\n", + " lr_labels.extend([lr] * n_seeds)\n", + "\n", + "df_1 = pd.concat(dfs_1, axis=1)\n", + "df_2 = pd.concat(dfs_2, axis=1)\n", + "df_3 = pd.concat(dfs_3, axis=1)\n", + "\n", + "params0 = load_params(SWEEP_DIR, LRS[0])\n", + "t_ref = 0.01 # seconds (t_ref=10 timesteps × 1 ms/timestep)\n", + "t_iter = t_ref * params0[\"batchsize\"] * params0[\"len_epoch\"]\n", + "t_max = (params0[\"n_epochs\"] + 1) * t_iter" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "MAX_T_ID = 64 * 30\n", + "num_epochs = 200\n", + "sec_per_epoch = 320\n", + "upper_axis_color = \"red\"\n", + "ts = np.arange(0, t_max, t_iter)\n", + "\n", + "fig, ax = plt.subplots(\n", + " 1, 3, figsize=(12 / 2.54, 5.5 / 2.54), sharey=True\n", + ") # Reduced figsize\n", + "\n", + "# 1. Plotting\n", + "lines = ax[0].plot(\n", + " ts[:MAX_T_ID], df_1.iloc[:MAX_T_ID, :], label=[f\"lr={lr}\" for lr in lr_labels]\n", + ")\n", + "ax[1].plot(ts[:MAX_T_ID], df_2.iloc[:MAX_T_ID, :])\n", + "ax[2].plot(ts[:MAX_T_ID], df_3.iloc[:MAX_T_ID, :])\n", + "\n", + "# Configure Left Panel Y-Axis\n", + "ax[0].set_ylabel(r\"$\\measuredangle (\\mathbf W^T, \\mathbf B)\\ [\\mathrm{deg}]$\")\n", + "ax[0].set_ylim(bottom=0.0)\n", + "\n", + "\n", + "# Define a custom formatter for scientific notation on the tick itself\n", + "# e.g. 100000 -> 1x10^5 or 10^5\n", + "def scientific_notation_formatter(x, pos):\n", + " if x == 0:\n", + " return \"0\"\n", + " # Format as 1e5, 2e5, etc.\n", + " s = \"{:.0e}\".format(x)\n", + " # Convert '1e+05' to latex '10^5' or '2 \\cdot 10^5'\n", + " base, exponent = s.split(\"e\")\n", + " exponent = int(exponent)\n", + " if base == \"1\":\n", + " return r\"$10^{%d}$\" % exponent\n", + " else:\n", + " return r\"$%s \\cdot 10^{%d}$\" % (base, exponent)\n", + "\n", + "\n", + "for i in range(3):\n", + " # Apply the custom formatter to the primary x-axis\n", + " ax[i].xaxis.set_major_formatter(ticker.FuncFormatter(scientific_notation_formatter))\n", + "\n", + " ax[i].minorticks_on()\n", + " ax[i].grid(which=\"major\")\n", + "\n", + " # --- SECONDARY AXIS ---\n", + " secax = ax[i].secondary_xaxis(\n", + " \"top\", functions=(lambda t: t / sec_per_epoch, lambda e: e * sec_per_epoch)\n", + " )\n", + " secax.set_xticks(np.arange(0, num_epochs + 1, 100))\n", + " ax[i].axvline(\n", + " num_epochs * sec_per_epoch,\n", + " linestyle=\"--\",\n", + " linewidth=0.8,\n", + " color=upper_axis_color,\n", + " )\n", + "\n", + " if i == 1:\n", + " secax.set_xlabel(\"Epochs\", color=upper_axis_color)\n", + "\n", + " secax.tick_params(axis=\"x\", colors=upper_axis_color)\n", + " secax.spines[\"top\"].set_color(upper_axis_color)\n", + "\n", + " # Internal Titles\n", + " ax[i].text(\n", + " 0.95,\n", + " 0.95,\n", + " f\"FC {i+1}\",\n", + " transform=ax[i].transAxes,\n", + " horizontalalignment=\"right\",\n", + " verticalalignment=\"top\",\n", + " bbox=dict(facecolor=\"white\", alpha=0.6, edgecolor=\"none\", pad=2),\n", + " )\n", + "\n", + "ax[1].set_xlabel(r\"$t$ [s]\")\n", + "\n", + "# --- FIGURE LEVEL LEGEND ---\n", + "handles, labels = ax[0].get_legend_handles_labels()\n", + "fig.legend(\n", + " handles,\n", + " labels,\n", + " loc=\"upper center\",\n", + " bbox_to_anchor=(0.5, 1.07),\n", + " ncol=len(labels),\n", + " frameon=True,\n", + ")\n", + "\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "fig.savefig(FIG_DIR / \"salnet_time.png\", dpi=300)\n", + "fig.savefig(FIG_DIR / \"salnet_time.pdf\")\n", + "fig.savefig(FIG_DIR / \"salnet_time.svg\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.10" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/scripts/symm_net/salnet_symm.py b/scripts/symm_net/salnet_symm.py index 30edb18..eca8eda 100644 --- a/scripts/symm_net/salnet_symm.py +++ b/scripts/symm_net/salnet_symm.py @@ -56,8 +56,13 @@ def parse_tags(s: str | None) -> list[str]: parser.add_argument( "--output-dir", type=str, default="../../results/symm_net", dest="output_dir" ) +parser.add_argument("--seed", type=int, default=0) +parser.add_argument("--run-dir", type=str, default=None, dest="run_dir") args = parser.parse_args() +torch.manual_seed(args.seed) +np.random.seed(args.seed) + params["n_epochs"] = args.n_epochs params["len_epoch"] = args.len_epoch params["batchsize"] = args.batchsize @@ -99,10 +104,14 @@ def append_metric(scalars: dict[str, Any], key: str, value: float) -> None: # Run setup # --------------------------- -run_dir = create_run_dirs(base_dir=args.output_dir, tags=tags) +if args.run_dir is not None: + run_dir = args.run_dir + os.makedirs(os.path.join(run_dir, "figs", "weights"), exist_ok=True) +else: + run_dir = create_run_dirs(base_dir=args.output_dir, tags=tags) metrics: dict[str, Any] = { - "params": params, + "params": {**params, "seed": args.seed}, "scalars": {}, } diff --git a/scripts/symm_net/sweep.py b/scripts/symm_net/sweep.py index 9ea1360..4f87db3 100644 --- a/scripts/symm_net/sweep.py +++ b/scripts/symm_net/sweep.py @@ -12,7 +12,7 @@ Subset example: python sweep.py --datasets cifar10 --algos bp sal --n-seeds 2 -Re-running is safe: finished runs (metrics.json present) are skipped. +Finished runs (metrics.json present) are skipped, if sweep is executed again. """ import argparse diff --git a/scripts/symm_net/sweep_creator_symm.py b/scripts/symm_net/sweep_creator_symm.py new file mode 100644 index 0000000..cc635be --- /dev/null +++ b/scripts/symm_net/sweep_creator_symm.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 + +"""Generate a jobs.sh file for SLURM array job submission. + +This is step 1 of the SLURM workflow for reproducing the SymmNet paper figure. +Each line of the output file contains one `python salnet_symm.py ...` command, +one per learning_rate × seed combination. Already-completed runs (those with a +metrics.json in the expected output directory) are skipped. + +Full workflow +------------- +1. Generate the job list:: + + python sweep_creator.py [--lrs ...] [--n-seeds N] + + This writes ``jobs.sh`` in the current directory. + +2. Submit as a SLURM array job (snapshots jobs.sh first for reproducibility):: + + bash slurm_submit.sh + + Internally, slurm_submit.sh calls ``sbatch --array=1-N slurm.sh``. + Each array task reads its own command from the snapshot and executes it. + +3. After all jobs finish, run ``plots.ipynb`` to reproduce the figure. + +Notes +----- +- ``jobs.sh`` must not exist before running this script (guard against + accidental overwrites); delete it manually to regenerate. +- To run locally instead of on SLURM, use ``sweep.py`` directly. +""" + +import argparse +from pathlib import Path + +DEFAULT_LRS = [0.01, 0.02, 0.04, 0.08] +PY_FILE = "salnet_symm.py" +JOBS_FILE = "jobs.sh" + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + p.add_argument( + "--lrs", + nargs="+", + type=float, + default=DEFAULT_LRS, + metavar="LR", + help="Learning rates to sweep (default: 0.01 0.02, 0.04, 0.08).", + ) + p.add_argument( + "--n-seeds", + type=int, + default=5, + dest="n_seeds", + help="Number of seeds per learning rate, numbered 0 … N-1 (default: 5).", + ) + p.add_argument( + "--n-workers", + type=int, + default=1, + dest="n_workers", + help="Number of parallel workers (default: 1 = sequential).", + ) + p.add_argument( + "--sweep-name", + default="puresymm", + dest="sweep_name", + help="Subdirectory name under --base-dir (default: puresymm).", + ) + p.add_argument( + "--base-dir", + default="../../results/symm_net/puresymm", + dest="base_dir", + help="Root output directory (default: ../../results/symm_net/puresymm).", + ) + # forwarded to salnet_symm.py + p.add_argument("--n-epochs", type=int, default=2000, dest="n_epochs") + p.add_argument("--len-epoch", type=int, default=100, dest="len_epoch") + p.add_argument("--batchsize", type=int, default=64) + return p.parse_args() + + +def main() -> None: + args = parse_args() + sweep_dir = Path(args.base_dir) / args.sweep_name + total = len(args.lrs) * args.n_seeds + done = 0 + + print(f"Sweep: {total} runs → {sweep_dir}") + print(f" lrs: {args.lrs}") + print(f" seeds: 0 … {args.n_seeds - 1}") + print(f" workers: {args.n_workers}") + print() + + if Path(JOBS_FILE).exists(): + print(f"{JOBS_FILE} already exists. Stop here!") + return + + proc_calls = [] + for lr in args.lrs: + for seed in range(args.n_seeds): + run_dir = sweep_dir / f"lr_{lr}" / f"seed_{seed}" + done += 1 + + if (run_dir / "metrics.json").exists(): + print(f"[{done}/{total}] Skip lr_{lr}/seed_{seed}") + continue + + print(f"[{done}/{total}] Run lr_{lr}/seed_{seed}") + proc_calls.append( + ( + f"python {PY_FILE} --lr {lr} --seed {seed} --n_epochs {args.n_epochs} " + f"--len_epoch {args.len_epoch} --batchsize {args.batchsize} " + f"--run-dir {run_dir}\n" + ) + ) + + with open(JOBS_FILE, "w") as f: + f.writelines(proc_calls) + + +if __name__ == "__main__": + main() diff --git a/scripts/symm_net/sweep_symm.py b/scripts/symm_net/sweep_symm.py new file mode 100644 index 0000000..9f4e0b4 --- /dev/null +++ b/scripts/symm_net/sweep_symm.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Sweep launcher for pure SAL symmetrization experiments (salnet_symm.py). + +Iterates over learning rates × seeds. Each combination is stored in a structured +directory tree that plot_puresymm.ipynb can read directly. + +Sequential run (default): + python sweep_symm.py + +Parallel run with 4 workers: + python sweep_symm.py --n-workers 4 + +Subset example: + python sweep_symm.py --lrs 0.001 0.01 --n-seeds 2 --n-epochs 10 + +Finished runs (metrics.json present) are skipped, if sweep is executed again. +""" + +import argparse +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +DEFAULT_LRS = [0.01, 0.02, 0.04, 0.08] + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + p.add_argument( + "--lrs", + nargs="+", + type=float, + default=DEFAULT_LRS, + metavar="LR", + help="Learning rates to sweep (default: 0.01 0.02, 0.04, 0.08).", + ) + p.add_argument( + "--n-seeds", + type=int, + default=5, + dest="n_seeds", + help="Number of seeds per learning rate, numbered 0 … N-1 (default: 5).", + ) + p.add_argument( + "--n-workers", + type=int, + default=1, + dest="n_workers", + help="Number of parallel workers (default: 1 = sequential).", + ) + p.add_argument( + "--sweep-name", + default="puresymm", + dest="sweep_name", + help="Subdirectory name under --base-dir (default: puresymm).", + ) + p.add_argument( + "--base-dir", + default="../../results/symm_net/puresymm", + dest="base_dir", + help="Root output directory (default: ../../results/symm_net/puresymm).", + ) + # forwarded to salnet_symm.py + p.add_argument("--n-epochs", type=int, default=2000, dest="n_epochs") + p.add_argument("--len-epoch", type=int, default=100, dest="len_epoch") + p.add_argument("--batchsize", type=int, default=64) + return p.parse_args() + + +def _run_one( + run_dir: Path, + lr: float, + seed: int, + n_epochs: int, + len_epoch: int, + batchsize: int, +) -> tuple[str, int | None]: + """Run a single salnet_symm.py job. + + Returns: + A (label, returncode) tuple. returncode is None if the run was skipped. + """ + label = f"lr={lr}/seed_{seed}" + if (run_dir / "metrics.json").exists(): + return label, None # already done + result = subprocess.run( + [ + sys.executable, + "salnet_symm.py", + "--seed", + str(seed), + "--run-dir", + str(run_dir), + "--lr", + str(lr), + "--n_epochs", + str(n_epochs), + "--len_epoch", + str(len_epoch), + "--batchsize", + str(batchsize), + ], + check=False, + ) + return label, result.returncode + + +def main() -> None: + args = parse_args() + sweep_dir = Path(args.base_dir) / args.sweep_name + total = len(args.lrs) * args.n_seeds + + print(f"Sweep: {total} runs → {sweep_dir}") + print(f" lrs: {args.lrs}") + print(f" seeds: 0 … {args.n_seeds - 1}") + print(f" workers: {args.n_workers}") + print() + + runs = [ + (sweep_dir / f"lr_{lr}" / f"seed_{seed}", lr, seed) + for lr in args.lrs + for seed in range(args.n_seeds) + ] + + done = 0 + with ThreadPoolExecutor(max_workers=args.n_workers) as pool: + futures = { + pool.submit( + _run_one, + run_dir, + lr, + seed, + args.n_epochs, + args.len_epoch, + args.batchsize, + ): i + for i, (run_dir, lr, seed) in enumerate(runs) + } + for future in as_completed(futures): + label, code = future.result() + done += 1 + if code is None: + print(f"[{done}/{total}] Skip {label}") + elif code != 0: + print(f"[{done}/{total}] FAIL {label} (exit {code})") + else: + print(f"[{done}/{total}] Done {label}") + + print(f"\nSweep complete. Results in {sweep_dir}") + + +if __name__ == "__main__": + main() From fa87f6c70daa899208cd26eb6e68c91894a909ba Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Thu, 16 Apr 2026 15:42:06 +0200 Subject: [PATCH 16/19] Add description to plotting notebooks --- mystyle.mpl | 1 + scripts/symm_net/plot_puresymm.ipynb | 7 ++++++- scripts/symm_net/plots.ipynb | 31 +++++++++++++++++++--------- 3 files changed, 28 insertions(+), 11 deletions(-) diff --git a/mystyle.mpl b/mystyle.mpl index 92e01e2..9c1c7aa 100644 --- a/mystyle.mpl +++ b/mystyle.mpl @@ -7,6 +7,7 @@ # use Latex: (uncomment the following two commands if latex is installed on system) # text.usetex: true # font.family: serif +# text.latex.preamble: \usepackage{amssymb} # text and font and labets etc. settings # axes.labelsize: 8 diff --git a/scripts/symm_net/plot_puresymm.ipynb b/scripts/symm_net/plot_puresymm.ipynb index 5c8131d..1982c21 100644 --- a/scripts/symm_net/plot_puresymm.ipynb +++ b/scripts/symm_net/plot_puresymm.ipynb @@ -5,7 +5,12 @@ "id": "0", "metadata": {}, "source": [ - "# Plots for pure symmetrization in SALNet" + "# Plots for pure symmetrization in SALNet (fig. 9)\n", + "\n", + "In this experiment, we investigate the convergence speed of SAL alone in the SALNet, when the forward weights are constant.\n", + "This helps us to determine the optimal tradeoff between fast convergence and low alignment angles by choosing an appropriate learning rate.\n", + "\n", + "The raw data is simulated by `salnet_symm.py` and stored in `../../results/symm_net/puresymm`. " ] }, { diff --git a/scripts/symm_net/plots.ipynb b/scripts/symm_net/plots.ipynb index 81b0ee3..8705c93 100644 --- a/scripts/symm_net/plots.ipynb +++ b/scripts/symm_net/plots.ipynb @@ -5,7 +5,18 @@ "id": "0", "metadata": {}, "source": [ - "# Plots for symmnet" + "# Plot the results of the SymmNet experiments (fig. 8)\n", + "\n", + "In this experiment, we apply SAL to the training of a deep convolutional neural network and compare its performance to competing symmetrization algorithms. \n", + "Learning happens in two phases: In each epoch, the forward weights $W$ are trained by standard gradient descent. Then, the trainable parameters (i.e., weights and biases) are copied over to a spiking network of equivalent architecture, where SAL (or RDD) is applied to the backward weights $B$.\n", + "\n", + "The forward training is done in a standard (non-spiking) ANN implemented in `pytorch`, where we use a custom-made backward function. We use a independent weight matrix $B$ in the computation of the backward path via autograd. During an epoch, $B$ is kept constant (similar to feedback alignment). \n", + "\n", + "Additionally, the implementation allows for other symmetrization schemes directly in the ANN such as feedback alignment, a variant of sign-concordant feedback alignment and the Kolen-Pollack algorithm.\n", + "\n", + "We test our setting on the CIFAR-10, Fashion MNIST and SVHN datasets.\n", + "\n", + "The training itself happens in `main_salnet.py`, which stores the simulated raw data in `../../results/symm_net/`. This notebook aggregates the raw data and plots it (see fig. 8 in the paper)." ] }, { @@ -37,8 +48,7 @@ "SAVEFIG = True\n", "\n", "# Point this at the sweep directory produced by sweep.py\n", - "# SWEEP_DIR = Path(\"../../results/symm_net/sweep\")\n", - "SWEEP_DIR = Path(\"/tmp/symmnet_test/sweep/\")\n", + "SWEEP_DIR = Path(\"../../results/symm_net/sweep\")\n", "\n", "FIG_DIR = Path(\"../../figs/symm_net\")\n", "FIG_DIR.mkdir(parents=True, exist_ok=True)" @@ -53,9 +63,7 @@ "source": [ "# define the style etc.\n", "mpl.style.use(\"../../mystyle.mpl\")\n", - "plt.style.use(\"tableau-colorblind10\")\n", - "plt.rcParams[\"text.usetex\"] = True\n", - "mpl.rcParams[\"text.latex.preamble\"] = r\"\\usepackage{amssymb}\"" + "plt.style.use(\"tableau-colorblind10\")" ] }, { @@ -347,7 +355,10 @@ " plot_epochs(fig, axes_upper[1], df, label)\n", "\n", "axes_upper[1].set_xlabel(\"epochs\")\n", - "axes_upper[1].set_ylabel(r\"$\\measuredangle (\\mathbf W^T, \\mathbf B)\\ [\\mathrm{deg}]$\")\n", + "# USE THE FOLLOWING LINE IF YOU HAVE LATEX INSTALLED:\n", + "# axes_upper[1].set_ylabel(r\"$\\measuredangle (\\mathbf W^T, \\mathbf B)\\ [\\mathrm{deg}]$\")\n", + "# If not, this one:\n", + "axes_upper[1].set_ylabel(r\"angle $( W^T, B)$ [deg]\")\n", "axes_upper[1].set_yticks([0, 45, 90])\n", "axes_upper[1].set_title(\"FC 1\")" ] @@ -564,9 +575,9 @@ "outputs": [], "source": [ "if SAVEFIG:\n", - " fig.savefig(\"../figs/salnet.png\", dpi=300)\n", - " fig.savefig(\"../figs/salnet.pdf\")\n", - " fig.savefig(\"../figs/salnet.svg\")" + " fig.savefig(FIG_DIR / \"salnet.png\", dpi=300)\n", + " fig.savefig(FIG_DIR / \"salnet.pdf\")\n", + " fig.savefig(FIG_DIR / \"salnet.svg\")" ] } ], From c55578b95c0ecd395a9ff6c8f99e01ffd433bda0 Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Thu, 16 Apr 2026 16:22:51 +0200 Subject: [PATCH 17/19] Update README with the explenation for the sweeps --- README.md | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 9ca7cdb..7ef9b5a 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,7 @@ To reproduce the data and plots for figure 7 (i.e., the weight scatter plots com #### Figure 8 A minimal working example for the deep learning experiment is provided by `scripts/symm_net/main_salnet.py` and the corresponding parameter file `exp_setting.yaml`. +As forward training of the convolutional neural networks is implemented in pytorch, we highly recommend to use a GPU. Usage: `python main_salnet.py -f exp_setting.yaml -s --dataset --tags ` @@ -99,9 +100,24 @@ Usage: `python main_salnet.py -f exp_setting.yaml -s --data - `dataset`: choose one the following datasets: `cifar10`, `svhn`, `mnist`, `fmnist` - optionally, you can pass a list of descriptive tags to keep tack of your runs. +The typical execution time of `main_salnet.py` amounts to ca. 15 minutes for the training modes that do not involve a spiking neural network, in SAL-mode ca 30-45 minutes and in RDD-mode ca. 2 hours. + +For conveniently reproducing any of the data shown in figure 8, we provide two workflows: +- `sweep.py`: suitable for small scale parameter sweeps. It directly launcher the required sub-processes. Example usage: `python sweep.py --datasets cifar10 --algos bp fa sal`. +See `python sweep.py --help` for all available settings. +- For large scale parameter sweeps (for instance to reproduce the 105 runs for all datasets, algorithms and seeds) on an HPC cluster (with SLURM), we provided the following workflow that needs only minimal adaptation to the available system. + 1. Run `sweep_creator.py`: It creates `jobs.sh` which contains all `main_salnet.py`-calls with the relevant settings. +See `python sweep_creator.py --help` for all available settings. + 2. Modify `slurm.sh` to specify the relevant settings for your HPC cluster. + 3. Run `bash slurm_submit.sh`. It will call `slurm.sh` internally and start a slurm array job. + 4. The data can be plotted with `scripts/symm_net/plots.ipynb` #### Figure 9 -A minimal working example for the Time evolution of SAL in the SALNet is provided by `scripts/symm_net/salnet_symm.py`. +A minimal working example for the time evolution of SAL in the SALNet is provided by `scripts/symm_net/salnet_symm.py`. + +Example Usage: `python salnet_symm.py --lr 0.01 --n_epochs 200 --len_epoch 500`. + +For conveniently reproducing the data shown in figure 9, we provide the same workflows as explained above with the files `sweep_symm.py` for quick small scale scans and `sweep_creator_symm.py` to launch all 20 runs as a slurm array job. #### Figure 10 From 3cf2dacfc0cd33ac1a106a6a13247c88b015391e Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Thu, 16 Apr 2026 16:48:31 +0200 Subject: [PATCH 18/19] symmnet package version bump --- symmnet/README.md | 11 ----------- symmnet/pyproject.toml | 2 +- symmnet/src/symmnet/__init__.py | 3 +++ 3 files changed, 4 insertions(+), 12 deletions(-) delete mode 100644 symmnet/README.md diff --git a/symmnet/README.md b/symmnet/README.md deleted file mode 100644 index 38550ba..0000000 --- a/symmnet/README.md +++ /dev/null @@ -1,11 +0,0 @@ - -# My adaptation of the "RDD net" used in Guerguiev et al, 2020 -to rate and compare the symmetrization capabilities of various spiking and non- -spiking algorithms in deep networks. - -Specifically, I implement SAL here and compare it against RDD... - - -## Get started: - -TODO diff --git a/symmnet/pyproject.toml b/symmnet/pyproject.toml index 1cf5ddb..4e5ff35 100644 --- a/symmnet/pyproject.toml +++ b/symmnet/pyproject.toml @@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta" [project] name = "symmnet" -version = "0.1" +version = "0.2" description = "Compare symmetriation algorithm in deep neural networks" readme = "README.md" authors = [{ name = "Timo Gierlich", email = "timo.gierlich@unibe.ch" }] diff --git a/symmnet/src/symmnet/__init__.py b/symmnet/src/symmnet/__init__.py index 50931f1..485b003 100644 --- a/symmnet/src/symmnet/__init__.py +++ b/symmnet/src/symmnet/__init__.py @@ -37,3 +37,6 @@ from .conv_net import ConvNet # noqa from .rdd_net import RDDNet # noqa +from .sal_net import SALNet # noqa + +__version__ = "0.2" From 2639e2f5271d7d4195bcf951d9cd16a8e6b6733a Mon Sep 17 00:00:00 2001 From: Timo Gierlich Date: Fri, 17 Apr 2026 12:38:20 +0200 Subject: [PATCH 19/19] BUGFIX: calculate means in plot_puresymm --- scripts/symm_net/plot_puresymm.ipynb | 27 ++++++++++++++++----------- scripts/symm_net/plots.ipynb | 7 +++---- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/scripts/symm_net/plot_puresymm.ipynb b/scripts/symm_net/plot_puresymm.ipynb index 1982c21..273a8b7 100644 --- a/scripts/symm_net/plot_puresymm.ipynb +++ b/scripts/symm_net/plot_puresymm.ipynb @@ -58,9 +58,7 @@ "source": [ "# define the style etc.\n", "mpl.style.use(\"../../mystyle.mpl\")\n", - "plt.style.use(\"tableau-colorblind10\")\n", - "plt.rcParams[\"text.usetex\"] = True\n", - "mpl.rcParams[\"text.latex.preamble\"] = r\"\\usepackage{amssymb}\"" + "plt.style.use(\"tableau-colorblind10\")" ] }, { @@ -91,7 +89,8 @@ " with open(path) as f:\n", " data = json.load(f)\n", " df[seed_dir.name] = data[\"scalars\"].get(metric_key, [])\n", - " return df\n", + " df_stats = pd.DataFrame({\"mean\": df.mean(axis=1), \"std\": df.std(axis=1)})\n", + " return df_stats\n", "\n", "\n", "def load_params(sweep_dir: Path, lr: float) -> dict:\n", @@ -112,17 +111,20 @@ "metadata": {}, "outputs": [], "source": [ - "dfs_1, dfs_2, dfs_3, lr_labels = [], [], [], []\n", + "dfs_1, dfs_2, dfs_3 = [], [], []\n", "for lr in LRS:\n", " n_seeds = len(list((SWEEP_DIR / f\"lr_{lr}\").glob(\"seed_*\")))\n", - " dfs_1.append(load_runs(SWEEP_DIR, lr, \"symm/angle/0\"))\n", - " dfs_2.append(load_runs(SWEEP_DIR, lr, \"symm/angle/1\"))\n", - " dfs_3.append(load_runs(SWEEP_DIR, lr, \"symm/angle/2\"))\n", - " lr_labels.extend([lr] * n_seeds)\n", + " dfs_1.append(load_runs(SWEEP_DIR, lr, \"symm/angle/0\")[\"mean\"])\n", + " dfs_2.append(load_runs(SWEEP_DIR, lr, \"symm/angle/1\")[\"mean\"])\n", + " dfs_3.append(load_runs(SWEEP_DIR, lr, \"symm/angle/2\")[\"mean\"])\n", "\n", "df_1 = pd.concat(dfs_1, axis=1)\n", + "df_1.columns = LRS\n", "df_2 = pd.concat(dfs_2, axis=1)\n", + "df_2.columns = LRS\n", "df_3 = pd.concat(dfs_3, axis=1)\n", + "df_3.columns = LRS\n", + "\n", "\n", "params0 = load_params(SWEEP_DIR, LRS[0])\n", "t_ref = 0.01 # seconds (t_ref=10 timesteps × 1 ms/timestep)\n", @@ -149,13 +151,16 @@ "\n", "# 1. Plotting\n", "lines = ax[0].plot(\n", - " ts[:MAX_T_ID], df_1.iloc[:MAX_T_ID, :], label=[f\"lr={lr}\" for lr in lr_labels]\n", + " ts[:MAX_T_ID], df_1.iloc[:MAX_T_ID, :], label=[f\"lr={lr}\" for lr in LRS]\n", ")\n", "ax[1].plot(ts[:MAX_T_ID], df_2.iloc[:MAX_T_ID, :])\n", "ax[2].plot(ts[:MAX_T_ID], df_3.iloc[:MAX_T_ID, :])\n", "\n", "# Configure Left Panel Y-Axis\n", + "# USE THE FOLLOWING LINE IF YOU HAVE LATEX INSTALLED:\n", "ax[0].set_ylabel(r\"$\\measuredangle (\\mathbf W^T, \\mathbf B)\\ [\\mathrm{deg}]$\")\n", + "# If not, this one:\n", + "ax[0].set_ylabel(r\"angle $( W^T, B)$ [deg]\")\n", "ax[0].set_ylim(bottom=0.0)\n", "\n", "\n", @@ -264,7 +269,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.10" + "version": "3.11.7" } }, "nbformat": 4, diff --git a/scripts/symm_net/plots.ipynb b/scripts/symm_net/plots.ipynb index 8705c93..735e075 100644 --- a/scripts/symm_net/plots.ipynb +++ b/scripts/symm_net/plots.ipynb @@ -28,7 +28,6 @@ "source": [ "import json\n", "import os\n", - "import pickle # noqa: F401 — kept for legacy compatibility if needed\n", "from pathlib import Path\n", "\n", "import numpy as np\n", @@ -583,9 +582,9 @@ ], "metadata": { "kernelspec": { - "display_name": "sal", + "display_name": "Python 3 (ipykernel)", "language": "python", - "name": "sal" + "name": "python3" }, "language_info": { "codemirror_mode": { @@ -597,7 +596,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.10" + "version": "3.11.7" } }, "nbformat": 4,