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 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/.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/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 diff --git a/scripts/symm_net/load_utils.py b/scripts/symm_net/load_utils.py index 7a9042c..0e25ccb 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,25 @@ 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, + int, + str | None, + str | None, +]: 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.") @@ -79,6 +100,14 @@ def settings_loader(): 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: @@ -119,4 +148,7 @@ def settings_loader(): 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 8578901..f8bacba 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 @@ -73,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"]) @@ -85,7 +97,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 +109,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] = [] @@ -114,18 +126,25 @@ def append_metric(metrics_dict, key, value): # 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: 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, "dataset": dataset.__name__, + "algo": section, + "seed": seed, "scalars": {}, } @@ -188,20 +207,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 +258,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 +293,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 +305,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 +350,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 +396,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 +597,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/plot_puresymm.ipynb b/scripts/symm_net/plot_puresymm.ipynb new file mode 100644 index 0000000..273a8b7 --- /dev/null +++ b/scripts/symm_net/plot_puresymm.ipynb @@ -0,0 +1,277 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# 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`. " + ] + }, + { + "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\")" + ] + }, + { + "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", + " 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", + " \"\"\"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 = [], [], []\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\")[\"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", + "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 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", + "# 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.11.7" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/scripts/symm_net/plots.ipynb b/scripts/symm_net/plots.ipynb new file mode 100644 index 0000000..735e075 --- /dev/null +++ b/scripts/symm_net/plots.ipynb @@ -0,0 +1,604 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# 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)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import os\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", + "\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\")" + ] + }, + { + "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", + "# 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\")" + ] + }, + { + "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(FIG_DIR / \"salnet.png\", dpi=300)\n", + " fig.savefig(FIG_DIR / \"salnet.pdf\")\n", + " fig.savefig(FIG_DIR / \"salnet.svg\")" + ] + } + ], + "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.11.7" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/scripts/symm_net/salnet_symm.py b/scripts/symm_net/salnet_symm.py index 493fdc9..eca8eda 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 [] @@ -55,8 +56,13 @@ def parse_tags(s): 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 @@ -70,7 +76,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 +87,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] = [] @@ -98,10 +104,14 @@ def append_metric(scalars, key, value): # 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 = { - "params": params, +metrics: dict[str, Any] = { + "params": {**params, "seed": args.seed}, "scalars": {}, } @@ -125,7 +135,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 +167,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 +204,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/scripts/symm_net/slurm.sh b/scripts/symm_net/slurm.sh new file mode 100644 index 0000000..8bb87f6 --- /dev/null +++ b/scripts/symm_net/slurm.sh @@ -0,0 +1,21 @@ +#!/bin/bash +# --- HPC cluster specific SLURM settings ------------------------ +# add your settings here. +#SBATCH --job-name="symmnet sweep" +#SBATCH --time= +#SBATCH ... +# ---------------------------------------------------------------- + +# --- HPC cluster specific setup:--------------------------------- +# e.g. load modules +# and activate python environment +# ---------------------------------------------------------------- + +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 new file mode 100644 index 0000000..963963a --- /dev/null +++ b/scripts/symm_net/slurm_submit.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# 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) +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.py b/scripts/symm_net/sweep.py new file mode 100644 index 0000000..4f87db3 --- /dev/null +++ b/scripts/symm_net/sweep.py @@ -0,0 +1,154 @@ +#!/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 + +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 + +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" + + +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( + "--n-workers", + type=int, + default=1, + dest="n_workers", + help="Number of parallel workers (default: 1 = sequential).", + ) + 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 _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 + + 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() + + 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}") + + +if __name__ == "__main__": + main() diff --git a/scripts/symm_net/sweep_creator.py b/scripts/symm_net/sweep_creator.py new file mode 100644 index 0000000..8a64b64 --- /dev/null +++ b/scripts/symm_net/sweep_creator.py @@ -0,0 +1,126 @@ +#!/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 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. + +Full workflow +------------- +1. Generate the job list:: + + python sweep_creator.py [--datasets ...] [--algos ...] [--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 + +ALL_DATASETS = ["cifar10", "fmnist", "svhn"] +ALL_ALGOS = ["bp", "fa", "bp_w_fa", "akrout", "scfa", "sal", "rdd"] +PARAM_FILE = "exp_settings.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} --seed {seed} --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_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() 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" 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..80f4c0d 100644 --- a/symmnet/src/symmnet/datasets.py +++ b/symmnet/src/symmnet/datasets.py @@ -1,10 +1,18 @@ #!/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, path="../datasets"): +def cifar10( + batch_size: int, path: str = _DEFAULT_DATA_DIR +) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 3 shape = (32, 32) # Data augmentation and normalization for training @@ -42,7 +50,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 = _DEFAULT_DATA_DIR +) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 1 shape = (28, 28) # Data augmentation and normalization for training @@ -80,7 +90,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 = _DEFAULT_DATA_DIR +) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 1 shape = (28, 28) # Data augmentation and normalization for training @@ -118,7 +130,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 = _DEFAULT_DATA_DIR +) -> tuple[DataLoader, DataLoader, int, tuple[int, int]]: n_channels = 3 shape = (160, 160) # Data augmentation and normalization for training @@ -156,7 +170,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 = _DEFAULT_DATA_DIR +) -> 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_original.py b/symmnet/src/symmnet/stdwi.py similarity index 89% rename from symmnet/src/symmnet/stdwi_original.py rename to symmnet/src/symmnet/stdwi.py index b46516c..43b6a7d 100644 --- a/symmnet/src/symmnet/stdwi_original.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]