diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 000000000..109e03581 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,8 @@ +.git +**/__pycache__ +*.pyc +data +logs +dist +build +.venv diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml index bf8650782..c2e70ed91 100644 --- a/.github/workflows/black.yml +++ b/.github/workflows/black.yml @@ -2,6 +2,9 @@ name: Black Formater on: [push, pull_request] +permissions: + contents: read + jobs: lint: runs-on: ubuntu-latest diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 000000000..6b248e8d6 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,51 @@ +name: Publish to PyPI + +on: + release: + types: [published] + workflow_dispatch: + +permissions: + contents: read + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: actions/setup-python@v7 + with: + python-version: "3.13" + - name: Build sdist and wheel + run: | + python -m pip install --upgrade build twine + python -m build + python -m twine check dist/* + - name: Check release tag matches pyproject version + if: github.event_name == 'release' + run: | + version=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml','rb'))['tool']['poetry']['version'])") + tag="${GITHUB_REF_NAME#v}" + if [ "$version" != "$tag" ]; then + echo "Release tag $GITHUB_REF_NAME does not match pyproject version $version" + exit 1 + fi + - uses: actions/upload-artifact@v7 + with: + name: dist + path: dist/ + + publish: + needs: build + runs-on: ubuntu-latest + environment: + name: pypi + url: https://pypi.org/project/bindsnet/ + permissions: + id-token: write + steps: + - uses: actions/download-artifact@v8 + with: + name: dist + path: dist/ + - uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml index 88936dbf4..18d293846 100644 --- a/.github/workflows/python-app.yml +++ b/.github/workflows/python-app.yml @@ -9,6 +9,9 @@ on: pull_request: branches: [ master ] +permissions: + contents: read + jobs: build: @@ -22,7 +25,7 @@ jobs: python-version: 3.13 - name: Install Poetry env: - POETRY_VERSION: 2.1.2 + POETRY_VERSION: 2.4.3 run: | curl -sSL https://install.python-poetry.org | python - -y &&\ poetry config virtualenvs.create false diff --git a/.github/workflows/pythonpackage.yml b/.github/workflows/pythonpackage.yml index bf44399b1..689080d42 100644 --- a/.github/workflows/pythonpackage.yml +++ b/.github/workflows/pythonpackage.yml @@ -2,6 +2,9 @@ name: Python package on: [push] +permissions: + contents: read + jobs: build: @@ -19,7 +22,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install Poetry env: - POETRY_VERSION: 2.1.2 + POETRY_VERSION: 2.4.3 run: | curl -sSL https://install.python-poetry.org | python - -y - name: Install dependencies diff --git a/.vscode/launch.json b/.vscode/launch.json index f08810a61..91a770239 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -9,6 +9,7 @@ "type": "python", "request": "launch", "program": "${file}", + "python": "/home/hananel/miniconda3/envs/bindsNET/bin/python", "console": "integratedTerminal", "justMyCode": false } diff --git a/CHANGELOG.md b/CHANGELOG.md index 056b7c4d6..5f254749e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,42 @@ see the [GitHub releases / tags](https://github.com/BindsNET/bindsnet/releases). ## [Unreleased] +## [0.3.4 (PyPI)] - 2026-09-16 + +First PyPI upload since 0.2.7. It is built from the `master` branch on this date, not +from the GitHub tag `0.3.4`, so `pip install bindsnet==0.3.4` contains everything in +this section **in addition to** the tag. The Zenodo archive +[10.5281/zenodo.20695116](https://doi.org/10.5281/zenodo.20695116) is the tag only. +Results can differ between the two: see the `MCC_learning.PostPre` entry under Changed. + +### Packaging +- Published to PyPI by `.github/workflows/publish.yml` (PyPI trusted publishing; runs + when a GitHub Release is published, or by hand). +- Removed install requirements that no module in `bindsnet/` or `examples/` imports: + `Cython`, `scikit-build`, `foolbox`, `numba`. +- `torch` is now `>=2.14,<3` and `torchvision` `>=0.29,<1` instead of exact pins; + `poetry.lock` still pins the tested versions (torch 2.14.0, torchvision 0.29.0). +- README: `pip install bindsnet`, PyPI badge, and absolute links and logo URL so the + PyPI project page renders. +- `pyproject.toml` package metadata moved from `[tool.poetry]` to the standard + `[project]` table (Poetry 2 deprecates the old one); `[tool.poetry]` now only routes + torch/torchvision to the CUDA 13.0 wheel index. `poetry.lock` resolves to the same + packages and versions. Build backend `poetry-core>=2.0`; the unused `setup.py` is removed. +- Poetry 2.4.3 in CI (was 2.1.2) and in `CONTRIBUTING.md` (said 1.1.8; `poetry shell` + replaced by `poetry env activate`, which Poetry 2 uses). +- `Dockerfile` rewritten. The old one could not build: its CUDA 11.1 base image, the + `get-poetry.py` installer and the `.python-version` file it copied no longer exist. + The new one uses `python:3.13-slim`, Poetry 2.4.3 and `poetry.lock`. The README no + longer links the Docker Hub image `hqkhan/bindsnet` (last updated 2019-01-28). +- Remaining links to the old `Hananel-Hazan/bindsnet` repository point to `BindsNET/bindsnet`. + +### Tests +- `test_perf_equivalence.py`: the batch-1 checks of the fused `addmm_` STDP update + required bit-for-bit equality with the un-fused formula. That holds on some CPUs + and not on GitHub's CI runners, where 5 tests failed. They now accept float32 + rounding (the tolerance already used for batch>1) and warn with the size of any + difference. + ### Added - Reproducibility/transparency docs: `DATA.md` (dataset & stimulus declaration), `REPRODUCING.md` (model→script→command→seed map), and a @@ -14,9 +50,101 @@ see the [GitHub releases / tags](https://github.com/BindsNET/bindsnet/releases). - `CHANGELOG.md`. - `examples/breakout/README.md` documenting the `trained_shallow_ANN.pt` provenance. +### Added +- `bindsnet.learning.DiehlAndCook` and `bindsnet.learning.MCC_learning.DiehlAndCook`: + the post-spike-only STDP of Diehl & Cook (2015), Sect. 2.3, + `dw = eta (x_pre - x_tar)(w_max - w)^mu` (keyword arguments `x_tar`, `mu`). + `DiehlAndCook2015(learning_rule=..., learning_rule_kwargs=...)` selects it; the + model's default stays `PostPre` so published results are unchanged. +- Multicompartment `Weight` features forward extra keyword arguments to their + learning rule. +- `bindsnet/learning/README.md`: rules, source papers, equation numbers, tests and + pitfalls (moved from the top-level README). + ### Changed +- `MCC_learning.PostPre` no longer multiplies its update by the simulation step + `dt`; like the classic `PostPre` and Morrison et al. (2008) eqs. 13-14 it is a + per-spike increment. Identical at `dt = 1`; at other steps the effective learning + rate is now `nu` instead of `nu * dt`. - README Python requirement aligned to `>=3.11,<3.14`; added a reproducible-install note. - `pyproject.toml` version bumped to 0.3.4 to match the released tag. +- Performance pass on the per-timestep hot paths (numerics unchanged; every item + is pinned by `test/network/test_perf_equivalence.py`, and a seeded old-vs-new + comparison of 59 networks was bit-identical except three batch>1 weight + matrices that differ by one float32 rounding step): + - `PostPre` / `Hebbian` on dense `Connection` and `MulticompartmentConnection` + apply the STDP update with one fused `addmm_` instead of materialising the + `[batch, source.n, target.n]` outer product (dense 784->1000 STDP, + batch 16, 250 steps on CPU: 9.6 s -> 0.42 s; Diehl & Cook 784->400, + batch 1: 1.6 s -> 0.26 s). + - `LearningRule.update` no longer multiplies the whole weight matrix by `1.0` + every step when no weight decay is configured. + - `LocalConnection1D/2D/3D` learning rules scale rows directly instead of + building an `[n, n]` identity matrix per step (64-filter local connection on + GPU: 61 MiB -> 2.8 MiB of per-step temporaries). `MSTDP`/`MSTDPET` keep + their post-synaptic trace as a `[batch, n, 1]` vector instead of a diagonal + matrix. + - `MSTDP` / `MSTDPET` cache `exp(-dt / tc)` and the default learning-rate + tensors instead of recomputing / re-copying them to the device each step. + - Neuron models update `v`, `refrac_count`, `theta`, `x`, ... in place with + the same operations in the same order, avoiding a `Module.__setattr__` + round-trip per assignment per step. + - `rank_order` encoding is vectorised. +- Benchmark script for the above: `examples/benchmark/hot_path_bench.py`. +- Learning rules validated against their source papers, with the equations cited in + `docs/source/models_spec.rst` and pinned by `test/network/test_learning_rule_specs.py`: + `PostPre` / `WeightDependentPostPre` / `Hebbian` against Morrison, Diesmann & + Gerstner (2008) eqs. (11)-(14); `MSTDP` / `MSTDPET` against Florian (2007) + eqs. (3.9)-(3.12) and (2.7)-(2.8) (equation numbers added to + `test_mstdp_florian.py`); `Rmax` against Vasilaki et al. (2009) eqs. (7), (8), (13). + The `MCC_learning` `PostPre` / `Hebbian` are checked to match the classic rules. +- Docstrings corrected: `Rmax` `tc_c` limits were stated backwards (`0` is the strict + policy-gradient rule, `inf` the naive Hebbian rule, Vasilaki et al. eq. 8); the + `MSTDP` / `MSTDPET` `zero_lag` comments called the un-lagged variant "exact Florian", + whereas the default one-step lag is Florian's discrete-time eq. (3.9). + +### Fixed +- `network.run(clamp=...)` / `unclamp` are now applied inside `Nodes.forward` before + the spike trace is updated, so a forced spike leaves a trace (and a suppressed one + does not). Previously the clamp was applied after the trace update, so clamped + spikes entered the same-step potentiation term of STDP rules but never the trace + used by later depression terms (affected `examples/mnist/supervised_mnist.py`). + Pinned by `TestClampEntersTraces` and the clamp-driven STDP window test. +- `network.to(device)` crashed on any `MulticompartmentConnection` (used by + `DiehlAndCook2015`) with `_apply() takes 2 positional arguments but 3 were + given`; `AbstractMulticompartmentConnection._apply` now accepts `recurse`. +- `Network.clone()` was broken outright: it called `torch.load` without + `weights_only=False`, so it raised `UnpicklingError` under PyTorch 2.6+, which + changed that default to `True`. It had no test and no caller in the tree, so the + breakage went unnoticed. Pinned by `TestNetwork.test_clone`. +- `Network.save()` called `torch.serialization.add_safe_globals([self])` with a + network instance where PyTorch expects a class. It did nothing useful and + corrupted PyTorch's safe-globals registry, so any later load in the same process + failed with `'Network' object has no attribute '__qualname__'`. Removed. Pinned by + `TestNetwork.test_clone_after_save`. +- `bindsnet.conversion.ann_to_snn` and `data_based_normalization` were broken when + given a path instead of a `torch.nn.Module`, for the same PyTorch 2.6 reason as + `Network.clone()`. Only the in-memory form was tested. Pinned by + `test_conversion_from_path` and `test_data_based_normalization_from_path`. + +### Security +- Documented that loading a saved network runs code. `bindsnet.network.load`, + `bindsnet.conversion.ann_to_snn` and `bindsnet.conversion.data_based_normalization` + read Python pickle files via `torch.load`, so a file from an untrusted source can + execute arbitrary code on load. This is the standard behaviour of `torch.load` + across the PyTorch ecosystem and is not a defect specific to BindsNET, but it was + undocumented. Added warnings to each function's docstring and a "Loading saved + networks and models" section to `SECURITY.md`. +- `bindsnet.network.load` gained a `weights_only` parameter, passed through to + `torch.load`. It defaults to `False`, which is required to read files written by + `Network.save` (those store the whole network object, not a tensor state dict), so + behaviour is unchanged. `weights_only=True` refuses code execution and is usable + only for files holding plain tensors. +- `SpokenMNIST` now reads its processed-data cache with `weights_only=True`. That + cache holds only tensors, so refusing code execution there costs nothing. Pinned by + `test/datasets/test_cache_serialization.py`. +- Reported by Gavin Branaa , who also prompted the three + `torch.load` fixes listed under Fixed above. Thank you. ## [0.3.4] - 2026-06-15 diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..539b1fab5 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,67 @@ +# BindsNET project notes + +## Python environment + +This project runs in the conda environment named `bindsNET`. + +- Interpreter: `/home/hananel/miniconda3/envs/bindsNET/bin/python` +- Activate: `conda activate bindsNET` +- The package is installed in editable mode from this directory, so edits to + `bindsnet/` are picked up without reinstalling. +- Torch 2.14 with CUDA 13.0 is installed there. Do not use the `base` env. + +Run tests with: + +```shell +conda run -n bindsNET python -m pytest -q +``` + +Run a single file with: + +```shell +conda run -n bindsNET python -m pytest -q test/network/test_learning.py +``` + +## Committing + +Use `sc "message"` instead of `git commit` (see the user's global instructions). + +## Learning rules: sources and how to validate + +Each rule is pinned to its paper's equations by a from-scratch reference test. +Read these before touching any rule: + +| Rule | Paper and equations | Test | +|---|---|---| +| `PostPre`, `WeightDependentPostPre`, `Hebbian` | Morrison, Diesmann & Gerstner 2008, *Biol. Cybern.* 98:459, eqs. 11-14 (traces Sect. 2.3) | `test/network/test_learning_rule_specs.py` | +| `DiehlAndCook` (classic and MCC) | Diehl & Cook 2015, *Front. Comput. Neurosci.* 9:99, Sect. 2.3: post-spike-only, dw = eta (x_pre - x_tar)(w_max - w)^mu | `test/network/test_learning_rule_specs.py` | +| `MSTDP`, `MSTDPET` | Florian 2007, *Neural Comput.* 19:1468, eqs. 3.9-3.12, 2.7-2.8 | `test/network/test_mstdp_florian.py` | +| `Rmax` | Vasilaki et al. 2009, *PLoS Comput. Biol.* 5:e1000586, eqs. 7, 8, 13 | `test/network/test_learning_rule_specs.py` | + +Facts that were wrong in docstrings once and are now fixed (do not reintroduce): +- `MSTDP` default `zero_lag=False` **is** Florian's discrete eq. 3.9 (reward at a + step multiplies the previous step's eligibility). `zero_lag=True` is the + un-lagged variant, not "exact Florian". +- `Rmax` `tc_c = 0` is strict policy gradient; `inf` is naive Hebbian. +- `clamp` spikes enter the trace (applied in `Nodes.forward` before the trace update). + +- Pair STDP rules carry no `dt` factor (the MCC `PostPre` used to; removed + 2026-09-06). `MSTDPET` keeps the paper's `dt` (Florian eq. 2.7). +- `PostPre` is not Diehl & Cook 2015's rule; the paper's rule is `DiehlAndCook`. + `DiehlAndCook2015` keeps `PostPre` as default (published replication); opt in + with `learning_rule=MCC_learning.DiehlAndCook`. + +User-facing summary of all this: `bindsnet/learning/README.md` (keep it in sync +with `docs/source/models_spec.rst`). + +## Performance changes: the rule + +Any change to a per-timestep path must ship with a test that pins it to the +formula it replaced (`test/network/test_perf_equivalence.py`). Before claiming +"no change in results", run the same seeded networks on the old code (a git +worktree of the previous commit) and the new code and compare with +`torch.equal`; only batch>1 summation-order differences (about 1e-7) are +acceptable, and must be stated. Exception: the batch-1 fused `addmm_` STDP update +matches the un-fused formula bit for bit on some CPUs but not on GitHub's CI runners +(max abs diff 1.5e-8 to 3.0e-8, CI run 35124573539, 2026-09-16), so those tests use +the same tolerance and warn with the size. Benchmark: `examples/benchmark/hot_path_bench.py`. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 1ca0f8aef..465c62cd5 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -3,12 +3,12 @@ To clone this project locally, issue ```shell -git clone https://github.com/Hananel-Hazan/bindsnet.git # clones bindsnet repository +git clone https://github.com/BindsNET/bindsnet.git # clones bindsnet repository ``` in the directory of your choice. This will place the repository's code in a directory titled `bindsnet`. -Install the project with [Poetry](https://python-poetry.org/) (current supported version - 1.1.8) +Install the project with [Poetry](https://python-poetry.org/) (version 2.4.3 is used in CI) ```shell poetry install @@ -16,7 +16,7 @@ poetry run pre-commit install ``` -Now you can access the project environment with `poetry shell` or run commands with `poetry run `. For example, `poetry run python examples/mnist/conv_mnist.py`. +Now you can activate the project environment with `eval $(poetry env activate)` or run commands with `poetry run `. For example, `poetry run python examples/mnist/conv_mnist.py`. Please make sure the `Poetry` environment is activated when you commit your files! The `git commit` command will invoke `pre-commit`, which is installed with Poetry too. IDEs like PyCharm have plugins for `Poetry` and will activate the environment automatically. @@ -29,6 +29,17 @@ poetry run pytest Notable changes are recorded in [`CHANGELOG.md`](CHANGELOG.md); please add an entry to the `Unreleased` section in your pull request. +## Releasing to PyPI + +1. Set `version` in `pyproject.toml` and `CITATION.cff`, and move the `Unreleased` + entries in `CHANGELOG.md` under the new version. +2. Merge to `master`, then publish a GitHub Release whose tag equals that version. +3. `.github/workflows/publish.yml` builds the package, checks that the tag matches the + `pyproject.toml` version, and uploads to PyPI through trusted publishing (no token). + It can also be started by hand from the Actions tab. + +A version number can be uploaded to PyPI only once, even after deletion. + All development should take place on a branch separate from master. To create a branch, issue ```shell @@ -102,6 +113,6 @@ where `[origin]` is the name of the remote repository, and `[branch-name]` is th __Note__: See [push.default](https://git-scm.com/docs/git-config#git-config-pushdefault) for more information. -To merge your changes into the `master` branch (the definitive version of the project's code), open a pull request on the [webpage](https://github.com/Hananel-Hazan/bindsnet) of the project. You can select the `base` branch (typically `master`, to merge changes _into_ the definitive version of the code) and the `compare` branch (say, `dan`, if I added a new feature locally and want to add it to the project code). You may add an optional extended description of your pull request changes. If there are merge conflicts at this stage, you may fix these using GitHub's pull request review interface. +To merge your changes into the `master` branch (the definitive version of the project's code), open a pull request on the [webpage](https://github.com/BindsNET/bindsnet) of the project. You can select the `base` branch (typically `master`, to merge changes _into_ the definitive version of the code) and the `compare` branch (say, `dan`, if I added a new feature locally and want to add it to the project code). You may add an optional extended description of your pull request changes. If there are merge conflicts at this stage, you may fix these using GitHub's pull request review interface. Assign reviewer(s) from the group of project contributors to perform a code review of your pull request. If the reviewer(s) are happy with your changes, you may then merge it in to the `master` branch. _Code review is crucial for the development of this project_, as the whole team should be held accountable for all changes. diff --git a/Dockerfile b/Dockerfile index 2a4dc25d3..b4446b9e3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,66 +1,33 @@ -ARG DEPS=development -ARG NVIDIA_30XX=false - -FROM nvidia/cuda:11.1-base AS base-default - -ARG DEBIAN_FRONTEND=noninteractive -RUN apt-get update && apt-get install --no-install-recommends --no-install-suggests -y \ - build-essential libgeos-dev liblzma-dev libssl-dev libbz2-dev curl vim python3.8-dev python-dev git libffi-dev \ - libglib2.0-0 libsm6 libxext6 libblas-dev libatlas-base-dev ffmpeg \ - && rm -rf /var/lib/apt/lists/* - -# install pyenv -ENV PYENV_ROOT=$HOME/.pyenv -ENV PATH $PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH -RUN curl -L https://github.com/pyenv/pyenv-installer/raw/master/bin/pyenv-installer | bash -RUN echo 'eval "$(pyenv init -)"' >> $HOME/.bashrc - -# install python version specified in the .python-version file -COPY .python-version . -RUN PYTHON_VERSION=$(cat .python-version) pyenv install $PYTHON_VERSION && pyenv global $PYTHON_VERSION && pyenv rehash - -# install poetry and our package -ENV POETRY_NO_INTERACTION=1\ - # send python output directory to stdout - PYTHONUNBUFFERED=1\ - PIP_NO_CACHE_DIR=off\ - PIP_DISABLE_PIP_VERSION_CHECK=on\ - PIP_DEFAULT_TIMEOUT=100\ - POETRY_HOME="/opt/poetry"\ - VENV_PATH="/opt/pysetup/.venv"\ - -# install poetry and our package -ENV POETRY_NO_INTERACTION=1 \ - # send python output directory to stdout - PYTHONUNBUFFERED=1 \ - PIP_NO_CACHE_DIR=off \ - PIP_DISABLE_PIP_VERSION_CHECK=on \ - PIP_DEFAULT_TIMEOUT=100 - -ENV PATH="$POETRY_HOME/bin:$VENV_PATH/bin:$PATH" POETRY_VERSION=1.1.8 - -RUN mkdir $HOME/opt/ && \ - curl -sSL https://raw.githubusercontent.com/python-poetry/poetry/master/get-poetry.py | python - &&\ - poetry config virtualenvs.create false +# BindsNET with the dependency set pinned in poetry.lock. +# docker build -t bindsnet . # with dev tools (pytest, black, ...) +# docker build --build-arg DEPS=main -t bindsnet . # runtime dependencies only +# docker run --rm -it --gpus all bindsnet bash # --gpus needs the NVIDIA container toolkit +# The torch wheels from the cu130 index bundle their CUDA libraries, so no CUDA base +# image is needed; the host only needs an NVIDIA driver. +FROM python:3.13-slim + +ARG DEPS=dev +ARG POETRY_VERSION=2.4.3 + +ENV PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + POETRY_NO_INTERACTION=1 \ + POETRY_VIRTUALENVS_CREATE=false + +# libgl1 and libglib2.0-0 are needed by opencv-python. +RUN apt-get update && apt-get install --no-install-recommends -y \ + git libgl1 libglib2.0-0 \ + && rm -rf /var/lib/apt/lists/* \ + && pip install "poetry==${POETRY_VERSION}" WORKDIR /bindsnet -RUN mkdir bindsnet && touch bindsnet/__init__.py ## empty package for Poetry to add to path +# Dependencies first, so editing the code does not re-download them. COPY pyproject.toml poetry.lock README.md ./ +RUN if [ "$DEPS" = "main" ]; then poetry install --no-root --only main; \ + else poetry install --no-root; fi \ + && rm -rf /root/.cache/pypoetry -FROM base-default AS base-production -RUN poetry install --no-dev # this will only install production dependencies - -FROM base-default AS base-development -RUN poetry install - -FROM base-${DEPS} AS nvidia-30xx-false -RUN rm -rf $HOME/.cache/pypoetry/artifacts # remove downloaded wheels - -# a fix for NVIDIA 30xx GPUs -FROM installed AS nvidia-30xx-true - -RUN python -m pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html - -FROM nvidia-30xx-${NVIDIA_30XX} AS final COPY . . +RUN poetry install --only-root diff --git a/README.md b/README.md index 16f225369..d812177d5 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -

+

A Python package used for simulating spiking neural networks (SNNs) on CPUs or GPUs using [PyTorch](http://pytorch.org/) `Tensor` functionality. @@ -9,6 +9,7 @@ This package is used as part of ongoing research on applying SNNs, machine learn Check out the [BindsNET examples](https://github.com/BindsNET/bindsnet/tree/master/examples) for a collection of experiments, functions for the analysis of results, plots of experiment outcomes, and more. Documentation for the package can be found [here](https://bindsnet-docs.readthedocs.io). +[![PyPI](https://img.shields.io/pypi/v/bindsnet.svg)](https://pypi.org/project/bindsnet/) [![Build Status](https://github.com/BindsNET/bindsnet/actions/workflows/python-app.yml/badge.svg?branch=master)](https://github.com/BindsNET/bindsnet/actions/workflows/python-app.yml) [![CodeQL](https://github.com/BindsNET/bindsnet/actions/workflows/github-code-scanning/codeql/badge.svg)](https://github.com/BindsNET/bindsnet/actions/workflows/github-code-scanning/codeql) [![Documentation Status](https://readthedocs.org/projects/bindsnet-docs/badge/?version=latest)](https://bindsnet-docs.readthedocs.io/?badge=latest) @@ -32,7 +33,13 @@ poetry install Alternatively, the provided `Dockerfile` builds the full pinned stack (see *Using Docker* below). ## Using Pip -To install the most recent stable release from the GitHub repository +To install the latest release from [PyPI](https://pypi.org/project/bindsnet/) + +``` +pip install bindsnet +``` + +To install the current development code from the GitHub repository ``` pip install git+https://github.com/BindsNET/bindsnet.git @@ -53,26 +60,22 @@ pip install -e . To install the packages necessary to interface with the [OpenAI gym RL environments library](https://github.com/openai/gym), follow their instructions for installing the packages needed to run the RL environments simulator (on Linux / MacOS). ### Using Docker -[Link](https://hub.docker.com/r/hqkhan/bindsnet/) to Docker repository. - -We also provide a Dockerfile in which BindsNET and all of its dependencies come installed in. Issue +The `Dockerfile` installs BindsNET with the dependency versions pinned in `poetry.lock`. +From the top level directory of this project, issue ``` -docker build . +docker build -t bindsnet . ``` -at the top level directory of this project to create a docker image. -To change the name of the newly built image, issue -``` -docker tag -``` - -To run a container and get a bash terminal inside it, issue +(add `--build-arg DEPS=main` to leave out the development tools). To get a bash +terminal inside a container, issue ``` -docker run -it bash +docker run --rm -it --gpus all bindsnet bash ``` +`--gpus all` needs the NVIDIA Container Toolkit; leave it out to run on CPU only. + ## Getting started To run a near-replication of the SNN from [this paper](https://www.frontiersin.org/articles/10.3389/fncom.2015.00099/full#), issue @@ -86,6 +89,13 @@ There are a number of optional command-line arguments which can be passed in, in A number of other examples are available in the `examples` directory that are meant to showcase BindsNET's functionality. Take a look, and let us know what you think! +## Learning rules and their sources + +Each weight-changing rule is validated against the equations of the paper it +implements. The table of rules, papers, equation numbers and tests, plus the pitfalls +(reward timing, traces, the Diehl & Cook rule versus `PostPre`), is in +[`bindsnet/learning/README.md`](https://github.com/BindsNET/bindsnet/blob/master/bindsnet/learning/README.md). + ## Running the tests Issue the following to run the tests: @@ -101,11 +111,11 @@ Some tests will fail if Open AI `gym` is not installed on your machine. BindsNET ships no third-party datasets; its loaders fetch them from upstream sources. Every dataset and synthetic stimulus used by the examples, benchmarks, and dataset loaders — with source, retrieval method, license pointer, and spike-encoding -preprocessing — is declared in [DATA.md](DATA.md). +preprocessing — is declared in [DATA.md](https://github.com/BindsNET/bindsnet/blob/master/DATA.md). ## Reproducing results -[REPRODUCING.md](REPRODUCING.md) maps each shipped model and published claim to its +[REPRODUCING.md](https://github.com/BindsNET/bindsnet/blob/master/REPRODUCING.md) maps each shipped model and published claim to its model class, example script, exact command, seed, and expected output (e.g. the Diehl & Cook 2015 MNIST replication via `examples/mnist/eth_mnist.py`, and the Hazan et al. 2018 scaling benchmark). @@ -130,7 +140,7 @@ We simulated a network with a population of n Poisson input neurons with firing Several packages, including BRIAN and PyNEST, allow the setting of certain global preferences; e.g., the number of CPU threads, the number of OpenMP processes, etc. We chose these settings for our benchmark study in an attempt to maximize each library's speed, but note that BindsNET requires no setting of such options. Our approach, inheriting the computational model of PyTorch, appears to make the best use of the available hardware, and therefore makes it simple for practitioners to get the best performance from their system with the least effort.

-BindsNET%20Benchmark +BindsNET%20Benchmark

All simulations run on Ubuntu 16.04 LTS with Intel(R) Xeon(R) CPU E5-2687W v3 @ 3.10GHz, 128Gb RAM @ 2133MHz, and two GeForce GTX TITAN X (GM200) GPUs. Python 3.6 is used in all cases. Clock time was recorded for each simulation run. @@ -162,8 +172,10 @@ The concept DOI below always resolves to the latest version: > BindsNET contributors. *BindsNET*. Zenodo. https://doi.org/10.5281/zenodo.20695115 (For the exact release used, cite its version DOI; e.g. v0.3.4 is -[10.5281/zenodo.20695116](https://doi.org/10.5281/zenodo.20695116).) A machine-readable -citation is provided in [`CITATION.cff`](CITATION.cff). +[10.5281/zenodo.20695116](https://doi.org/10.5281/zenodo.20695116). That archive is the +GitHub tag `0.3.4`; the PyPI package `bindsnet==0.3.4` was built later and also +contains the changes listed in `CHANGELOG.md` under 0.3.4 (PyPI).) A machine-readable +citation is provided in [`CITATION.cff`](https://github.com/BindsNET/bindsnet/blob/master/CITATION.cff). ## Contributors diff --git a/SECURITY.md b/SECURITY.md index b730ff5fb..dad47acef 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -16,6 +16,41 @@ you were on. We aim to acknowledge reports within a few days. Security fixes are applied to the `master` branch and to the most recent release on PyPI. Older releases are not patched. +## Loading saved networks and models + +Saved network files are Python pickle files. Loading one runs whatever code is +stored inside it. This means: + +**Only load network or model files you created yourself, or that came from a +source you trust.** A file from a download, a model-sharing site, shared group +storage, or a collaborator can run arbitrary commands on your machine as soon +as you load it. There is no way to inspect a pickle file safely first. + +This affects the following functions, all of which load a file path you give +them: + +- `bindsnet.network.load` +- `bindsnet.conversion.ann_to_snn`, when passed a path instead of a module +- `bindsnet.conversion.data_based_normalization`, when passed a path instead of + a module + +This is the standard behaviour of `torch.load`, and it applies to saved models +across the PyTorch ecosystem rather than being specific to BindsNET. We do not +treat it as a vulnerability in BindsNET, because there is no way to load a +saved network object without it. We do treat it as something you need to be +told about clearly, which is what this section and the warnings in those +functions' documentation are for. + +`bindsnet.network.load` accepts `weights_only=True`, which asks PyTorch to +refuse to execute code while loading. It cannot read files written by +`Network.save`, because those store the whole network object rather than a +plain tensor state dictionary, so it is only useful for files you know contain +plain tensors. + +If you need to share a trained network with people who cannot verify where it +came from, share the weights as a tensor state dictionary rather than a pickled +network object. + ## Repository integrity BindsNET is a research library, and its git history is part of what users rely diff --git a/bindsnet/conversion/conversion.py b/bindsnet/conversion/conversion.py index dbcf7785f..3febb0f2b 100644 --- a/bindsnet/conversion/conversion.py +++ b/bindsnet/conversion/conversion.py @@ -88,6 +88,12 @@ def data_based_normalization( Use a dataset to rescale ANN weights and biases such that that the max ReLU activation is less than 1. + .. warning:: + When ``ann`` is a path, the file is loaded with :func:`torch.load` + using ``weights_only=False``, which runs code stored in the file. Only + pass a path to a file you created yourself or otherwise trust. Pass an + already-loaded ``torch.nn.Module`` to avoid this. + :param ann: Artificial neural network implemented in PyTorch. Accepts either ``torch.nn.Module`` or path to network saved using ``torch.save()``. @@ -99,7 +105,9 @@ def data_based_normalization( according to activations on the dataset. """ if isinstance(ann, str): - ann = torch.load(ann) + # weights_only=False is required to load a whole nn.Module. PyTorch 2.6 + # changed the default to True, which cannot read these files. + ann = torch.load(ann, weights_only=False) assert isinstance(ann, nn.Module) @@ -277,6 +285,12 @@ def ann_to_snn( Converts an artificial neural network (ANN) written as a ``torch.nn.Module`` into a near-equivalent spiking neural network. + .. warning:: + When ``ann`` is a path, the file is loaded with :func:`torch.load` + using ``weights_only=False``, which runs code stored in the file. Only + pass a path to a file you created yourself or otherwise trust. Pass an + already-loaded ``torch.nn.Module`` to avoid this. + :param ann: Artificial neural network implemented in PyTorch. Accepts either ``torch.nn.Module`` or path to network saved using ``torch.save()``. @@ -290,7 +304,9 @@ def ann_to_snn( :return: Spiking neural network implemented in PyTorch. """ if isinstance(ann, str): - ann = torch.load(ann) + # weights_only=False is required to load a whole nn.Module. PyTorch 2.6 + # changed the default to True, which cannot read these files. + ann = torch.load(ann, weights_only=False) else: ann = deepcopy(ann) diff --git a/bindsnet/datasets/spoken_mnist.py b/bindsnet/datasets/spoken_mnist.py index b8c6f544c..ce60a64d7 100644 --- a/bindsnet/datasets/spoken_mnist.py +++ b/bindsnet/datasets/spoken_mnist.py @@ -115,7 +115,10 @@ def _get_train(self, split: float = 0.8) -> Tuple[torch.Tensor, torch.Tensor]: else: # Load image data from disk if it has already been processed. print("Loading training data from serialized object file.\n") - audio, labels = torch.load(open(path, "rb")) + # weights_only=True is safe here: this cache file holds only + # tensors, written by torch.save above, and refusing code + # execution costs nothing. + audio, labels = torch.load(open(path, "rb"), weights_only=True) labels = torch.Tensor(labels) @@ -163,7 +166,10 @@ def _get_test(self, split: float = 0.8) -> Tuple[torch.Tensor, List[torch.Tensor else: # Load image data from disk if it has already been processed. print("Loading test data from serialized object file.\n") - audio, labels = torch.load(open(path, "rb")) + # weights_only=True is safe here: this cache file holds only + # tensors, written by torch.save above, and refusing code + # execution costs nothing. + audio, labels = torch.load(open(path, "rb"), weights_only=True) labels = torch.Tensor(labels) diff --git a/bindsnet/encoding/encodings.py b/bindsnet/encoding/encodings.py index d17f5dc32..7e299eed3 100644 --- a/bindsnet/encoding/encodings.py +++ b/bindsnet/encoding/encodings.py @@ -182,10 +182,11 @@ def rank_order( times *= time / times.max() # Extended through simulation time. times = torch.ceil(times).long() - # Create spike times tensor. + # Create spike times tensor (one spike per neuron whose time lies in + # ``(0, time)``; vectorised form of the per-neuron loop). spikes = torch.zeros(time, size, device=device).byte() - for i in range(size): - if 0 < times[i] < time: - spikes[times[i] - 1, i] = 1 + fire = (times > 0) & (times < time) + idx = fire.nonzero(as_tuple=False).squeeze(1) + spikes[(times[fire] - 1).to(device), idx.to(device)] = 1 return spikes.reshape(time, *shape) diff --git a/bindsnet/learning/MCC_learning.py b/bindsnet/learning/MCC_learning.py index ef20d44d1..a6ef5538c 100644 --- a/bindsnet/learning/MCC_learning.py +++ b/bindsnet/learning/MCC_learning.py @@ -13,6 +13,35 @@ from ..utils import im2col_indices +def _dense_outer_update_ok(rule, w: torch.Tensor) -> bool: + # language=rst + """ + Whether a rule's ``[batch, source.n] x [batch, target.n]`` outer-product + update can be applied with a single fused ``w.addmm_`` call: dense float32 + weights and one of the two default batch reductions (``torch.squeeze`` for + batch size 1, ``torch.sum`` otherwise), both equal to the matrix product's + contraction over the batch dimension. + """ + return ( + isinstance(w, torch.Tensor) + and not w.is_sparse + and w.dtype == torch.float32 + and rule.reduction in (torch.squeeze, torch.sum) + ) + + +def _batch_outer(a: torch.Tensor, b: torch.Tensor, reduction) -> torch.Tensor: + # language=rst + """ + ``reduction(bmm(a[:, :, None], b[:, None, :]), dim=0)``: the batch of outer + products of ``a`` (``[batch, n]``) and ``b`` (``[batch, m]``) reduced over the + batch. For the two default reductions this is the single matmul ``a^T b``. + """ + if reduction in (torch.squeeze, torch.sum): + return a.t() @ b + return reduction(torch.bmm(a.unsqueeze(2), b.unsqueeze(1)), dim=0) + + class MCC_LearningRule(ABC): # language=rst """ @@ -151,7 +180,9 @@ class PostPre(MCC_LearningRule): # language=rst """ Simple STDP rule involving both pre- and post-synaptic spiking activity. By default, - pre-synaptic update is negative and the post-synaptic update is positive. + pre-synaptic update is negative and the post-synaptic update is positive. Same + equations as ``bindsnet.learning.PostPre`` (Morrison et al. 2008, eqs. 13-14): the + update is a per-spike increment and does not depend on the simulation step ``dt``. """ def __init__( @@ -229,6 +260,25 @@ def _connection_update(self, **kwargs) -> None: class. """ batch_size = self.source.batch_size + w = self.feature_value + + if self.average_update == 0 and _dense_outer_update_ok(self, w): + # Fused path: ``w += alpha * s^T @ x`` folds the outer product, the + # batch reduction and the in-place update into one ``addmm_`` call, + # so no ``[batch, source.n, target.n]`` temporary is allocated. The + # learning rate is applied to the ``[batch, n]`` factor first (as in + # the un-fused formula) and spikes are exactly 0/1, so for batch + # size 1 the result is bit-identical. + if self.nu[0]: + source_s = self.source.s.view(batch_size, -1).float() + target_x = self.target.x.view(batch_size, -1) * self.nu[0] + w.addmm_(source_s.t(), target_x, alpha=-1.0) + if self.nu[1]: + target_s = self.target.s.view(batch_size, -1).float() * self.nu[1] + source_x = self.source.x.view(batch_size, -1) + w.addmm_(source_x.t(), target_s, alpha=1.0) + super().update() + return # Pre-synaptic update. if self.nu[0]: @@ -245,22 +295,15 @@ def _connection_update(self, **kwargs) -> None: ) % self.average_update if self.continues_update: - self.feature_value -= ( - torch.mean(self.average_buffer_pre, dim=0) * self.connection.dt - ) + self.feature_value -= torch.mean(self.average_buffer_pre, dim=0) elif self.average_buffer_index_pre == 0: - self.feature_value -= ( - torch.mean(self.average_buffer_pre, dim=0) * self.connection.dt - ) + self.feature_value -= torch.mean(self.average_buffer_pre, dim=0) else: if self.feature_value.is_sparse: - self.feature_value -= ( - torch.bmm(source_s, target_x) * self.connection.dt - ).to_sparse() + self.feature_value -= (torch.bmm(source_s, target_x)).to_sparse() else: - self.feature_value -= ( - self.reduction(torch.bmm(source_s, target_x), dim=0) - * self.connection.dt + self.feature_value -= self.reduction( + torch.bmm(source_s, target_x), dim=0 ) del source_s, target_x @@ -281,22 +324,15 @@ def _connection_update(self, **kwargs) -> None: ) % self.average_update if self.continues_update: - self.feature_value += ( - torch.mean(self.average_buffer_post, dim=0) * self.connection.dt - ) + self.feature_value += torch.mean(self.average_buffer_post, dim=0) elif self.average_buffer_index_post == 0: - self.feature_value += ( - torch.mean(self.average_buffer_post, dim=0) * self.connection.dt - ) + self.feature_value += torch.mean(self.average_buffer_post, dim=0) else: if self.feature_value.is_sparse: - self.feature_value += ( - torch.bmm(source_x, target_s) * self.connection.dt - ).to_sparse() + self.feature_value += (torch.bmm(source_x, target_s)).to_sparse() else: - self.feature_value += ( - self.reduction(torch.bmm(source_x, target_s), dim=0) - * self.connection.dt + self.feature_value += self.reduction( + torch.bmm(source_x, target_s), dim=0 ) del source_x, target_s @@ -380,18 +416,27 @@ def _connection_update(self, **kwargs) -> None: batch_size = self.source.batch_size - source_s = self.source.s.view(batch_size, -1).unsqueeze(2).float() - source_x = self.source.x.view(batch_size, -1).unsqueeze(2) - target_s = self.target.s.view(batch_size, -1).unsqueeze(1).float() - target_x = self.target.x.view(batch_size, -1).unsqueeze(1) + if _dense_outer_update_ok(self, self.feature_value): + # Fused path (see ``PostPre._connection_update``). + source_s = self.source.s.view(batch_size, -1).float() + source_x = self.source.x.view(batch_size, -1) + target_s = self.target.s.view(batch_size, -1).float() + target_x = self.target.x.view(batch_size, -1) + self.feature_value.addmm_(source_s.t(), target_x * self.nu[0], alpha=1.0) + self.feature_value.addmm_(source_x.t(), target_s * self.nu[1], alpha=1.0) + else: + source_s = self.source.s.view(batch_size, -1).unsqueeze(2).float() + source_x = self.source.x.view(batch_size, -1).unsqueeze(2) + target_s = self.target.s.view(batch_size, -1).unsqueeze(1).float() + target_x = self.target.x.view(batch_size, -1).unsqueeze(1) - # Pre-synaptic update. - update = self.reduction(torch.bmm(source_s, target_x), dim=0) - self.feature_value += self.nu[0] * update + # Pre-synaptic update. + update = self.reduction(torch.bmm(source_s, target_x), dim=0) + self.feature_value += self.nu[0] * update - # Post-synaptic update. - update = self.reduction(torch.bmm(source_x, target_s), dim=0) - self.feature_value += self.nu[1] * update + # Post-synaptic update. + update = self.reduction(torch.bmm(source_x, target_s), dim=0) + self.feature_value += self.nu[1] * update # Add polarities back to feature after updates if self.enforce_polarity: @@ -407,6 +452,109 @@ def reset_state_variables(self) -> None: """ +class DiehlAndCook(MCC_LearningRule): + # language=rst + """ + Post-synaptic-spike-only STDP of `Diehl & Cook (2015) + `_, Sect. 2.3: + + .. math:: + + \\Delta w = \\eta\\,(x_\\text{pre} - x_\\text{tar})\\,(w_\\max - w)^\\mu + + applied on every post-synaptic spike, with :math:`x_\\text{pre}` the source + spike trace (the paper's trace adds 1 per spike: use ``traces_additive=True``), + :math:`x_\\text{tar}` the target trace value, :math:`w_\\max` the upper end of + ``range``, and :math:`\\mu` the weight-dependence exponent. Pre-synaptic spikes do + not change the weight. Multicompartment twin of + ``bindsnet.learning.DiehlAndCook``. + """ + + def __init__( + self, + connection: AbstractMulticompartmentConnection, + feature_value: Union[torch.Tensor, float, int], + range: Optional[Sequence[float]] = None, + nu: Optional[Union[float, Sequence[float]]] = None, + reduction: Optional[callable] = None, + decay: float = 0.0, + **kwargs, + ) -> None: + # language=rst + """ + Constructor for the ``DiehlAndCook`` learning rule. + + :param connection: A ``MulticompartmentConnection`` whose weight feature the + rule will modify. + :param feature_value: The weight tensor. + :param range: ``[w_min, w_max]``; ``w_max`` must be finite. + :param nu: Learning rate :math:`\\eta`. A pair is accepted for API symmetry; + only the second (post-synaptic) entry is used. + :param reduction: Method for reducing parameter updates along the batch + dimension. + :param decay: Coefficient controlling rate of decay of the weights each iteration. + + Keyword arguments: + + :param float x_tar: Target pre-synaptic trace :math:`x_\\text{tar}` (default 0). + :param float mu: Weight-dependence exponent :math:`\\mu` (default 1). + """ + super().__init__( + connection=connection, + feature_value=feature_value, + range=[0.0, 1.0] if range is None else range, + nu=nu, + reduction=reduction, + decay=decay, + **kwargs, + ) + + assert self.source.traces, "Pre-synaptic nodes must record spike traces." + assert self.max is not None and np.isfinite( + self.max + ), "DiehlAndCook needs a finite upper weight bound (range[1] = w_max)." + + if isinstance(connection, MulticompartmentConnection): + self.update = self._connection_update + else: + raise NotImplementedError( + "This learning rule is not supported for this Connection type." + ) + + self.x_tar = float(kwargs.get("x_tar", 0.0)) + self.mu = float(kwargs.get("mu", 1.0)) + + def _connection_update(self, **kwargs) -> None: + # language=rst + """ + ``w += eta * ((x_pre - x_tar) outer s_post) * (w_max - w) ** mu`` reduced + over the batch. + """ + if not self.nu[1]: + super().update() + return + batch_size = self.source.batch_size + w = self.feature_value + source_x = self.source.x.view(batch_size, -1) - self.x_tar + target_s = self.target.s.view(batch_size, -1).float() + outer = _batch_outer(source_x, target_s, self.reduction) + factor = self.max - w + if self.mu != 1.0: + factor = factor.clamp(min=0.0) ** self.mu + if w.is_sparse: + w += (self.nu[1] * outer * factor.to_dense()).to_sparse() + else: + w += self.nu[1] * outer * factor + super().update() + + def reset_state_variables(self) -> None: + # language=rst + """ + Nothing to reset: the rule holds no state between steps, deriving each + update from the source trace, the target spikes and the current weight. + """ + + class MSTDP(MCC_LearningRule): # language=rst """ diff --git a/bindsnet/learning/README.md b/bindsnet/learning/README.md new file mode 100644 index 000000000..53cd6c7ff --- /dev/null +++ b/bindsnet/learning/README.md @@ -0,0 +1,73 @@ +# Learning rules and their sources + +Every weight-changing rule in `bindsnet.learning` is validated step by step against +the equations of the paper it implements. The equations are written out with their +numbers in [`docs/source/models_spec.rst`](../../docs/source/models_spec.rst) +("Learning rules") and pinned by the tests in the last column, which drive the rule +with the spike trains the network actually produces and compare the weights on every +step against a from-scratch transcription of the paper. + +| Rule (`bindsnet.learning`) | Multicompartment twin (`bindsnet.learning.MCC_learning`) | Paper | Test | +|---|---|---|---| +| `PostPre` | `PostPre` | Morrison, Diesmann & Gerstner (2008), *Biol. Cybern.* 98:459, eqs. 11-14: additive pair STDP with traces | `test/network/test_learning_rule_specs.py` | +| `WeightDependentPostPre` | - | same, soft-bounded: F+ = nu_post (w_max - w), F- = nu_pre (w - w_min) | same | +| `Hebbian` | `Hebbian` | BindsNET's own definition: both trace terms potentiate | same | +| `DiehlAndCook` | `DiehlAndCook` | Diehl & Cook (2015), *Front. Comput. Neurosci.* 9:99, Sect. 2.3: post-spike-only rule dw = eta (x_pre - x_tar)(w_max - w)^mu | same | +| `MSTDP`, `MSTDPET` | `MSTDP`, `MSTDPET` | Florian (2007), *Neural Comput.* 19:1468, eqs. 3.9-3.12 and 2.7-2.8 | `test/network/test_mstdp_florian.py` | +| `Rmax` | - | Vasilaki et al. (2009), *PLoS Comput. Biol.* 5:e1000586, eqs. 7, 8, 13 | `test/network/test_learning_rule_specs.py` | + +## Things that are easy to get wrong + +- **Reward timing.** The reward passed to `network.run` at a step multiplies the + eligibility built from the *previous* step's spikes. That is Florian's discrete + equation 3.9 and the default (`zero_lag=False`). `zero_lag=True` is the un-lagged + variant. +- **`PostPre` is not the Diehl & Cook rule.** `PostPre` is standard pair STDP: pre + spikes depress by the post trace, post spikes potentiate by the pre trace. The + `DiehlAndCook2015` model uses it by default because that is what the published + BindsNET replication used. The paper's own rule is `DiehlAndCook` (below). +- **`Rmax` `tc_c`.** `tc_c = 0` is the strict policy-gradient rule; `tc_c = inf` is + naive Hebbian (Vasilaki et al. eq. 8). +- **Traces.** `traces_additive=False` (default) resets the trace to `trace_scale` on + each spike (Morrison's saturating trace with A = 1); `traces_additive=True` adds + `trace_scale` per spike (the accumulating trace, which is what Diehl & Cook use). +- **Clamped spikes count.** Spikes forced with `network.run(clamp=...)` enter the + spike trace and the learning rules like any other spike; `unclamp` removes a + spike before the trace sees it. +- **No `dt` factor in pair STDP.** `PostPre`, `WeightDependentPostPre`, `Hebbian` + and `DiehlAndCook` are per-spike increments; the same spike pair changes the + weight by the same amount at any simulation step. (`MSTDPET` does carry the + paper's `dt` factor, eq. 2.7.) + +## Using the Diehl & Cook (2015) rule + +```python +from bindsnet.learning import DiehlAndCook +from bindsnet.network.topology import Connection + +conn = Connection(source, target, nu=(0.0, 1e-2), update_rule=DiehlAndCook, + wmin=0.0, wmax=1.0, x_tar=0.4, mu=1.0) +``` + +`source` must record traces; use `traces_additive=True` on the source layer for the +paper's accumulating trace. Only the second learning rate (post-synaptic) is used. +`wmax` must be finite. The paper does not give numbers for `x_tar` and `mu`; the +defaults (0 and 1) are BindsNET's, so set them for your experiment. + +With the `DiehlAndCook2015` model: + +```python +from bindsnet.learning.MCC_learning import DiehlAndCook +from bindsnet.models import DiehlAndCook2015 + +net = DiehlAndCook2015(n_inpt=784, n_neurons=100, inpt_shape=(1, 28, 28), + learning_rule=DiehlAndCook, + learning_rule_kwargs={"x_tar": 0.4, "mu": 1.0}) +``` + +This also switches the input layer to accumulating traces. Leaving `learning_rule` +unset keeps the published `PostPre` behaviour, so results in `REPRODUCING.md` are +unchanged. + +For any multicompartment `Weight` feature, extra keyword arguments (`x_tar`, `mu`, +`tc_plus`, ...) are forwarded to its learning rule. diff --git a/bindsnet/learning/__init__.py b/bindsnet/learning/__init__.py index 5a733783a..ef5a4bf23 100644 --- a/bindsnet/learning/__init__.py +++ b/bindsnet/learning/__init__.py @@ -1,6 +1,7 @@ from bindsnet.learning.learning import ( MSTDP, MSTDPET, + DiehlAndCook, Hebbian, LearningRule, NoOp, @@ -15,6 +16,7 @@ "PostPre", "WeightDependentPostPre", "Hebbian", + "DiehlAndCook", "MSTDP", "MSTDPET", "Rmax", diff --git a/bindsnet/learning/learning.py b/bindsnet/learning/learning.py index a0a8f027d..cf45ed9a8 100644 --- a/bindsnet/learning/learning.py +++ b/bindsnet/learning/learning.py @@ -51,6 +51,82 @@ def _conv_point_eligibility(connection, p_plus, p_minus, source_s, target_s, dim ) +def _dense_outer_update_ok(rule, w: torch.Tensor) -> bool: + # language=rst + """ + Whether a rule's ``[batch, source.n] x [batch, target.n]`` outer-product + update can be applied with a single fused ``w.addmm_`` call. + + This requires a dense float32 weight matrix, scalar learning rates (a + per-synapse ``nu`` tensor cannot be folded into ``alpha``), and one of the + two default batch reductions (``torch.squeeze`` for batch size 1, + ``torch.sum`` otherwise), both of which equal the matrix product's + contraction over the batch dimension. + """ + return ( + not w.is_sparse + and w.dtype == torch.float32 + and rule.nu.dim() == 1 + and rule.reduction in (torch.squeeze, torch.sum) + ) + + +def _row_scale(vec: torch.Tensor, mat: torch.Tensor) -> torch.Tensor: + # language=rst + """ + ``torch.bmm(torch.diag_embed(vec), mat)`` without building the + ``[batch, n, n]`` diagonal matrix: scales row ``i`` of ``mat`` by + ``vec[:, i]``. Bit-identical for finite values (the dropped terms are exact + zeros). + + :param vec: Tensor of shape ``[batch, n, 1]`` or ``[batch, n]``. + :param mat: Tensor of shape ``[batch, n, k]``. + """ + if vec.dim() == 2: + vec = vec.unsqueeze(2) + return vec * mat + + +def _cached_decay(rule, name: str) -> torch.Tensor: + # language=rst + """ + ``exp(-dt / rule.)`` computed once per ``(name, dt)`` and reused on + every timestep instead of re-evaluating the exponential each call. + """ + dt = float(rule.connection.dt) + tc = getattr(rule, name) + cache = rule.__dict__.setdefault("_decay_cache", {}) + # Keyed on the time constant's value too, so a later change to it is honoured. + key = (name, dt, float(tc)) + if key not in cache: + cache[key] = torch.exp(-dt / tc) + return cache[key] + + +def _reward_rates(rule, kwargs: dict, device: torch.device): + # language=rst + """ + Return the ``(a_plus, a_minus)`` learning-rate tensors for the reward-modulated + rules. The defaults (``1.0`` / ``-1.0``) are allocated once on ``device`` and + reused, so a run that does not pass ``a_plus``/``a_minus`` no longer copies a + fresh scalar to the device on every timestep. + """ + cache = rule.__dict__.setdefault("_rate_cache", {}) + if cache.get("device") != device: + cache["device"] = device + cache["a_plus"] = torch.tensor(1.0, device=device) + cache["a_minus"] = torch.tensor(-1.0, device=device) + a_plus = kwargs.get("a_plus", None) + a_minus = kwargs.get("a_minus", None) + a_plus = ( + cache["a_plus"] if a_plus is None else torch.as_tensor(a_plus, device=device) + ) + a_minus = ( + cache["a_minus"] if a_minus is None else torch.as_tensor(a_minus, device=device) + ) + return a_plus, a_minus + + class LearningRule(ABC): # language=rst """ @@ -118,8 +194,9 @@ def update(self) -> None: """ Abstract method for a learning rule update. """ - # Implement weight decay. - if self.weight_decay: + # Implement weight decay (1.0 is the no-decay default; skip the + # full-matrix multiply in that case). + if self.weight_decay != 1.0: self.connection.w *= self.weight_decay # Bound weights. @@ -249,9 +326,6 @@ def _local_connection1d_update(self, **kwargs) -> None: height_out = self.connection.conv_size target_x = self.target.x.reshape(batch_size, out_channels * height_out, 1) - target_x = target_x * torch.eye(out_channels * height_out).to( - self.connection.w.device - ) source_s = ( self.source.s.type(torch.float) .unfold(-1, kernel_height, stride) @@ -263,9 +337,6 @@ def _local_connection1d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out, 1 ) - target_s = target_s * torch.eye(out_channels * height_out).to( - self.connection.w.device - ) source_x = ( self.source.x.unfold(-1, kernel_height, stride) .reshape(batch_size, height_out, in_channels * kernel_height) @@ -275,11 +346,11 @@ def _local_connection1d_update(self, **kwargs) -> None: # Pre-synaptic update. if self.nu[0].any(): - pre = self.reduction(torch.bmm(target_x, source_s), dim=0) + pre = self.reduction(_row_scale(target_x, source_s), dim=0) self.connection.w -= self.nu[0] * pre.view(self.connection.w.size()) # Post-synaptic update. if self.nu[1].any(): - post = self.reduction(torch.bmm(target_s, source_x), dim=0) + post = self.reduction(_row_scale(target_s, source_x), dim=0) self.connection.w += self.nu[1] * post.view(self.connection.w.size()) super().update() @@ -303,9 +374,6 @@ def _local_connection2d_update(self, **kwargs) -> None: target_x = self.target.x.reshape( batch_size, out_channels * height_out * width_out, 1 ) - target_x = target_x * torch.eye(out_channels * height_out * width_out).to( - self.connection.w.device - ) source_s = ( self.source.s.type(torch.float) .unfold(-2, kernel_height, stride[0]) @@ -322,9 +390,6 @@ def _local_connection2d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out * width_out, 1 ) - target_s = target_s * torch.eye(out_channels * height_out * width_out).to( - self.connection.w.device - ) source_x = ( self.source.x.unfold(-2, kernel_height, stride[0]) .unfold(-2, kernel_width, stride[1]) @@ -339,11 +404,11 @@ def _local_connection2d_update(self, **kwargs) -> None: # Pre-synaptic update. if self.nu[0].any(): - pre = self.reduction(torch.bmm(target_x, source_s), dim=0) + pre = self.reduction(_row_scale(target_x, source_s), dim=0) self.connection.w -= self.nu[0] * pre.view(self.connection.w.size()) # Post-synaptic update. if self.nu[1].any(): - post = self.reduction(torch.bmm(target_s, source_x), dim=0) + post = self.reduction(_row_scale(target_s, source_x), dim=0) self.connection.w += self.nu[1] * post.view(self.connection.w.size()) super().update() @@ -369,9 +434,6 @@ def _local_connection3d_update(self, **kwargs) -> None: target_x = self.target.x.reshape( batch_size, out_channels * height_out * width_out * depth_out, 1 ) - target_x = target_x * torch.eye( - out_channels * height_out * width_out * depth_out - ).to(self.connection.w.device) source_s = ( self.source.s.type(torch.float) .unfold(-3, kernel_height, stride[0]) @@ -389,9 +451,6 @@ def _local_connection3d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out * width_out * depth_out, 1 ) - target_s = target_s * torch.eye( - out_channels * height_out * width_out * depth_out - ).to(self.connection.w.device) source_x = ( self.source.x.unfold(-3, kernel_height, stride[0]) .unfold(-3, kernel_width, stride[1]) @@ -407,11 +466,11 @@ def _local_connection3d_update(self, **kwargs) -> None: # Pre-synaptic update. if self.nu[0].any(): - pre = self.reduction(torch.bmm(target_x, source_s), dim=0) + pre = self.reduction(_row_scale(target_x, source_s), dim=0) self.connection.w -= self.nu[0] * pre.view(self.connection.w.size()) # Post-synaptic update. if self.nu[1].any(): - post = self.reduction(torch.bmm(target_s, source_x), dim=0) + post = self.reduction(_row_scale(target_s, source_x), dim=0) self.connection.w += self.nu[1] * post.view(self.connection.w.size()) super().update() @@ -423,6 +482,24 @@ def _connection_update(self, **kwargs) -> None: class. """ batch_size = self.source.batch_size + w = self.connection.w + + if _dense_outer_update_ok(self, w): + # Fused path: ``w += alpha * s^T @ x`` folds the outer product, the + # batch reduction and the in-place update into one ``addmm_`` call, + # so no ``[batch, source.n, target.n]`` temporary is allocated. + # Spikes are exactly 0/1, so the products are bit-identical to the + # un-fused formula for batch size 1. + if self.nu[0].any(): + source_s = self.source.s.view(batch_size, -1).float() + target_x = self.target.x.view(batch_size, -1) * self.nu[0] + w.addmm_(source_s.t(), target_x, alpha=-1.0) + if self.nu[1].any(): + target_s = self.target.s.view(batch_size, -1).float() * self.nu[1] + source_x = self.source.x.view(batch_size, -1) + w.addmm_(source_x.t(), target_s, alpha=1.0) + super().update() + return # Pre-synaptic update. if self.nu[0].any(): @@ -696,9 +773,6 @@ def _local_connection1d_update(self, **kwargs) -> None: height_out = self.connection.conv_size target_x = self.target.x.reshape(batch_size, out_channels * height_out, 1) - target_x = target_x * torch.eye(out_channels * height_out).to( - self.connection.w.device - ) source_s = ( self.source.s.type(torch.float) .unfold(-1, kernel_height, stride) @@ -710,9 +784,6 @@ def _local_connection1d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out, 1 ) - target_s = target_s * torch.eye(out_channels * height_out).to( - self.connection.w.device - ) source_x = ( self.source.x.unfold(-1, kernel_height, stride) .reshape(batch_size, height_out, in_channels * kernel_height) @@ -724,7 +795,7 @@ def _local_connection1d_update(self, **kwargs) -> None: # Pre-synaptic update. if self.nu[0].any(): - pre = self.reduction(torch.bmm(target_x, source_s), dim=0) + pre = self.reduction(_row_scale(target_x, source_s), dim=0) update -= ( self.nu[0] * pre.view(self.connection.w.size()) @@ -732,7 +803,7 @@ def _local_connection1d_update(self, **kwargs) -> None: ) # Post-synaptic update. if self.nu[1].any(): - post = self.reduction(torch.bmm(target_s, source_x), dim=0) + post = self.reduction(_row_scale(target_s, source_x), dim=0) update += ( self.nu[1] * post.view(self.connection.w.size()) @@ -762,9 +833,6 @@ def _local_connection2d_update(self, **kwargs) -> None: target_x = self.target.x.reshape( batch_size, out_channels * height_out * width_out, 1 ) - target_x = target_x * torch.eye(out_channels * height_out * width_out).to( - self.connection.w.device - ) source_s = ( self.source.s.type(torch.float) .unfold(-2, kernel_height, stride[0]) @@ -781,9 +849,6 @@ def _local_connection2d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out * width_out, 1 ) - target_s = target_s * torch.eye(out_channels * height_out * width_out).to( - self.connection.w.device - ) source_x = ( self.source.x.unfold(-2, kernel_height, stride[0]) .unfold(-2, kernel_width, stride[1]) @@ -800,7 +865,7 @@ def _local_connection2d_update(self, **kwargs) -> None: # Pre-synaptic update. if self.nu[0].any(): - pre = self.reduction(torch.bmm(target_x, source_s), dim=0) + pre = self.reduction(_row_scale(target_x, source_s), dim=0) update -= ( self.nu[0] * pre.view(self.connection.w.size()) @@ -808,7 +873,7 @@ def _local_connection2d_update(self, **kwargs) -> None: ) # Post-synaptic update. if self.nu[1].any(): - post = self.reduction(torch.bmm(target_s, source_x), dim=0) + post = self.reduction(_row_scale(target_s, source_x), dim=0) update += ( self.nu[1] * post.view(self.connection.w.size()) @@ -840,9 +905,6 @@ def _local_connection3d_update(self, **kwargs) -> None: target_x = self.target.x.reshape( batch_size, out_channels * height_out * width_out * depth_out, 1 ) - target_x = target_x * torch.eye( - out_channels * height_out * width_out * depth_out - ).to(self.connection.w.device) source_s = ( self.source.s.type(torch.float) .unfold(-3, kernel_height, stride[0]) @@ -860,9 +922,6 @@ def _local_connection3d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out * width_out * depth_out, 1 ) - target_s = target_s * torch.eye( - out_channels * height_out * width_out * depth_out - ).to(self.connection.w.device) source_x = ( self.source.x.unfold(-3, kernel_height, stride[0]) .unfold(-3, kernel_width, stride[1]) @@ -880,7 +939,7 @@ def _local_connection3d_update(self, **kwargs) -> None: # Pre-synaptic update. if self.nu[0].any(): - pre = self.reduction(torch.bmm(target_x, source_s), dim=0) + pre = self.reduction(_row_scale(target_x, source_s), dim=0) update -= ( self.nu[0] * pre.view(self.connection.w.size()) @@ -888,7 +947,7 @@ def _local_connection3d_update(self, **kwargs) -> None: ) # Post-synaptic update. if self.nu[1].any(): - post = self.reduction(torch.bmm(target_s, source_x), dim=0) + post = self.reduction(_row_scale(target_s, source_x), dim=0) update += ( self.nu[1] * post.view(self.connection.w.size()) @@ -1143,6 +1202,20 @@ def _connection_update(self, **kwargs) -> None: class. """ batch_size = self.source.batch_size + w = self.connection.w + + if _dense_outer_update_ok(self, w): + # Fused path (see ``PostPre._connection_update``): the learning rate + # is applied to the [batch, n] factor first, so the per-synapse + # product matches the un-fused ``nu * (s * x)`` bit for bit. + source_s = self.source.s.view(batch_size, -1).float() + source_x = self.source.x.view(batch_size, -1) + target_s = self.target.s.view(batch_size, -1).float() + target_x = self.target.x.view(batch_size, -1) + w.addmm_(source_s.t(), target_x * self.nu[0], alpha=1.0) + w.addmm_(source_x.t(), target_s * self.nu[1], alpha=1.0) + super().update() + return source_s = self.source.s.view(batch_size, -1).unsqueeze(2).float() source_x = self.source.x.view(batch_size, -1).unsqueeze(2) @@ -1178,9 +1251,6 @@ def _local_connection1d_update(self, **kwargs) -> None: height_out = self.connection.conv_size target_x = self.target.x.reshape(batch_size, out_channels * height_out, 1) - target_x = target_x * torch.eye(out_channels * height_out).to( - self.connection.w.device - ) source_s = ( self.source.s.type(torch.float) .unfold(-1, kernel_height, stride) @@ -1192,9 +1262,6 @@ def _local_connection1d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out, 1 ) - target_s = target_s * torch.eye(out_channels * height_out).to( - self.connection.w.device - ) source_x = ( self.source.x.unfold(-1, kernel_height, stride) .reshape(batch_size, height_out, in_channels * kernel_height) @@ -1203,11 +1270,11 @@ def _local_connection1d_update(self, **kwargs) -> None: ) # Pre-synaptic update. - pre = self.reduction(torch.bmm(target_x, source_s), dim=0) + pre = self.reduction(_row_scale(target_x, source_s), dim=0) self.connection.w += self.nu[0] * pre.view(self.connection.w.size()) # Post-synaptic update. - post = self.reduction(torch.bmm(target_s, source_x), dim=0) + post = self.reduction(_row_scale(target_s, source_x), dim=0) self.connection.w += self.nu[1] * post.view(self.connection.w.size()) super().update() @@ -1231,9 +1298,6 @@ def _local_connection2d_update(self, **kwargs) -> None: target_x = self.target.x.reshape( batch_size, out_channels * height_out * width_out, 1 ) - target_x = target_x * torch.eye(out_channels * height_out * width_out).to( - self.connection.w.device - ) source_s = ( self.source.s.type(torch.float) .unfold(-2, kernel_height, stride[0]) @@ -1250,9 +1314,6 @@ def _local_connection2d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out * width_out, 1 ) - target_s = target_s * torch.eye(out_channels * height_out * width_out).to( - self.connection.w.device - ) source_x = ( self.source.x.unfold(-2, kernel_height, stride[0]) .unfold(-2, kernel_width, stride[1]) @@ -1266,11 +1327,11 @@ def _local_connection2d_update(self, **kwargs) -> None: ) # Pre-synaptic update. - pre = self.reduction(torch.bmm(target_x, source_s), dim=0) + pre = self.reduction(_row_scale(target_x, source_s), dim=0) self.connection.w += self.nu[0] * pre.view(self.connection.w.size()) # Post-synaptic update. - post = self.reduction(torch.bmm(target_s, source_x), dim=0) + post = self.reduction(_row_scale(target_s, source_x), dim=0) self.connection.w += self.nu[1] * post.view(self.connection.w.size()) super().update() @@ -1296,9 +1357,6 @@ def _local_connection3d_update(self, **kwargs) -> None: target_x = self.target.x.reshape( batch_size, out_channels * height_out * width_out * depth_out, 1 ) - target_x = target_x * torch.eye( - out_channels * height_out * width_out * depth_out - ).to(self.connection.w.device) source_s = ( self.source.s.type(torch.float) .unfold(-3, kernel_height, stride[0]) @@ -1316,9 +1374,6 @@ def _local_connection3d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out * width_out * depth_out, 1 ) - target_s = target_s * torch.eye( - out_channels * height_out * width_out * depth_out - ).to(self.connection.w.device) source_x = ( self.source.x.unfold(-3, kernel_height, stride[0]) .unfold(-3, kernel_width, stride[1]) @@ -1333,11 +1388,11 @@ def _local_connection3d_update(self, **kwargs) -> None: ) # Pre-synaptic update. - pre = self.reduction(torch.bmm(target_x, source_s), dim=0) + pre = self.reduction(_row_scale(target_x, source_s), dim=0) self.connection.w += self.nu[0] * pre.view(self.connection.w.size()) # Post-synaptic update. - post = self.reduction(torch.bmm(target_s, source_x), dim=0) + post = self.reduction(_row_scale(target_s, source_x), dim=0) self.connection.w += self.nu[1] * post.view(self.connection.w.size()) super().update() @@ -1467,6 +1522,103 @@ def _conv3d_connection_update(self, **kwargs) -> None: super().update() +class DiehlAndCook(LearningRule): + # language=rst + """ + Post-synaptic-spike-only STDP of `Diehl & Cook (2015) + `_, Sect. 2.3: + + .. math:: + + \\Delta w = \\eta\\,(x_\\text{pre} - x_\\text{tar})\\,(w_\\max - w)^\\mu + + applied on every post-synaptic spike. :math:`x_\\text{pre}` is the source layer's + spike trace (the paper's trace adds 1 per spike, i.e. ``traces_additive=True``); + :math:`x_\\text{tar}` is the target trace value ("the higher the target value, the + lower the synaptic weight will be"); :math:`w_\\max` is the connection's ``wmax``; + :math:`\\mu` sets the weight dependence. Pre-synaptic spikes do not change the + weight, unlike ``PostPre``. The paper gives no numeric values for + :math:`x_\\text{tar}` and :math:`\\mu`; the defaults here (0 and 1) are BindsNET's. + """ + + def __init__( + self, + connection: AbstractConnection, + nu: Optional[Union[float, Sequence[float], Sequence[torch.Tensor]]] = None, + reduction: Optional[callable] = None, + weight_decay: float = 0.0, + **kwargs, + ) -> None: + # language=rst + """ + Constructor for the ``DiehlAndCook`` learning rule. + + :param connection: A ``Connection`` or ``LocalConnection`` whose weights the + rule will modify. It must have a finite ``wmax``. + :param nu: Learning rate :math:`\\eta`. A pair is accepted for API symmetry + with the other rules; only the second (post-synaptic) entry is used. + :param reduction: Method for reducing parameter updates along the batch + dimension. + :param weight_decay: Coefficient controlling rate of decay of the weights each + iteration. + + Keyword arguments: + + :param float x_tar: Target pre-synaptic trace :math:`x_\\text{tar}` (default 0). + :param float mu: Weight-dependence exponent :math:`\\mu` (default 1). + """ + super().__init__( + connection=connection, + nu=nu, + reduction=reduction, + weight_decay=weight_decay, + **kwargs, + ) + + assert self.source.traces, "Pre-synaptic nodes must record spike traces." + assert ( + connection.wmax != np.inf + ).all(), "DiehlAndCook needs a finite wmax (the paper's w_max)." + + if isinstance(connection, (Connection, LocalConnection)): + self.update = self._connection_update + else: + raise NotImplementedError( + "This learning rule is not supported for this Connection type." + ) + + self.x_tar = float(kwargs.get("x_tar", 0.0)) + self.mu = float(kwargs.get("mu", 1.0)) + + def _connection_update(self, **kwargs) -> None: + # language=rst + """ + ``w += eta * ((x_pre - x_tar) outer s_post) * (wmax - w) ** mu`` reduced + over the batch. + """ + if not self.nu[1].any(): + super().update() + return + batch_size = self.source.batch_size + w = self.connection.w + source_x = self.source.x.view(batch_size, -1) - self.x_tar + target_s = self.target.s.view(batch_size, -1).float() + if self.reduction in (torch.squeeze, torch.sum): + outer = source_x.t() @ target_s + else: + outer = self.reduction( + torch.bmm(source_x.unsqueeze(2), target_s.unsqueeze(1)), dim=0 + ) + factor = self.connection.wmax - w + if self.mu != 1.0: + factor = factor.clamp(min=0.0) ** self.mu + update = self.nu[1] * outer * factor + if w.is_sparse: + update = update.to_sparse() + self.connection.w += update + super().update() + + class MSTDP(LearningRule): # language=rst """ @@ -1529,10 +1681,13 @@ def __init__( self.tc_plus = torch.tensor(kwargs.get("tc_plus", 20.0)) self.tc_minus = torch.tensor(kwargs.get("tc_minus", 20.0)) - # If True, the reward at step t modulates the eligibility that already - # includes the spikes at step t (exact Florian 2007 timing). If False - # (default, backward-compatible) the eligibility is applied with a - # one-timestep lag. Currently honoured by the ``Connection`` path. + # Timing of reward vs. eligibility. Default (False) is Florian (2007) + # eq. 3.9, w(t+dt) = w(t) + gamma r(t+dt) zeta(t): the reward supplied + # at a step multiplies the eligibility built from the previous step's + # spikes. ``zero_lag=True`` instead multiplies the reward by the + # eligibility that already includes this step's spikes (a direct + # discretisation of the continuous-time eq. 3.4). Currently honoured by + # the ``Connection`` path. self.zero_lag = kwargs.get("zero_lag", False) def _connection_update(self, **kwargs) -> None: @@ -1575,31 +1730,35 @@ def _connection_update(self, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = kwargs.get("a_plus", 1.0) - if isinstance(a_plus, dict): - for k, v in a_plus.items(): - a_plus[k] = torch.tensor(v, device=self.connection.w.device) - else: - a_plus = torch.tensor(a_plus, device=self.connection.w.device) - a_minus = kwargs.get("a_minus", -1.0) - if isinstance(a_minus, dict): - for k, v in a_minus.items(): - a_minus[k] = torch.tensor(v, device=self.connection.w.device) + a_plus = kwargs.get("a_plus", None) + a_minus = kwargs.get("a_minus", None) + if isinstance(a_plus, dict) or isinstance(a_minus, dict): + if isinstance(a_plus, dict): + for k, v in a_plus.items(): + a_plus[k] = torch.tensor(v, device=self.connection.w.device) + elif a_plus is None: + a_plus = 1.0 + if isinstance(a_minus, dict): + for k, v in a_minus.items(): + a_minus[k] = torch.tensor(v, device=self.connection.w.device) + elif a_minus is None: + a_minus = -1.0 else: - a_minus = torch.tensor(a_minus, device=self.connection.w.device) + a_plus, a_minus = _reward_rates(self, kwargs, self.connection.w.device) # Update P^+/P^- traces and the point eligibility for this timestep. def _update_traces_and_eligibility(): - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= _cached_decay(self, "tc_plus") self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= _cached_decay(self, "tc_minus") self.p_minus += a_minus * target_s self.eligibility = torch.bmm( self.p_plus.unsqueeze(2), target_s.unsqueeze(1) ) + torch.bmm(source_s.unsqueeze(2), self.p_minus.unsqueeze(1)) # With zero_lag, fold in the current spikes before applying the update, - # so reward(t) multiplies eligibility(t) exactly as in Florian 2007. + # so reward(t) multiplies eligibility(t) (un-lagged variant; the default + # below is Florian 2007 eq. 3.9). if self.zero_lag: _update_traces_and_eligibility() @@ -1609,8 +1768,8 @@ def _update_traces_and_eligibility(): update = update.to_sparse() self.connection.w += self.nu[0] * update - # Default (backward-compatible): the eligibility computed here is applied - # on the next timestep (one-step lag). + # Default: the eligibility computed here is applied on the next + # timestep, as in Florian (2007) eq. 3.9 / eqs. 2.7-2.8. if not self.zero_lag: _update_traces_and_eligibility() @@ -1639,12 +1798,7 @@ def _local_connection1d_update(self, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = torch.tensor( - kwargs.get("a_plus", 1.0), device=self.connection.w.device - ) - a_minus = torch.tensor( - kwargs.get("a_minus", -1.0), device=self.connection.w.device - ) + a_plus, a_minus = _reward_rates(self, kwargs, self.connection.w.device) # Compute weight update based on the eligibility value of the past timestep. update = reward * self.eligibility @@ -1669,9 +1823,6 @@ def _local_connection1d_update(self, **kwargs) -> None: self.p_minus = self.p_minus.reshape( batch_size, out_channels * height_out, 1 ) - self.p_minus = self.p_minus * torch.eye(out_channels * height_out).to( - self.connection.w.device - ) # Reshaping spike occurrences. source_s = ( @@ -1685,18 +1836,15 @@ def _local_connection1d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out, 1 ) - target_s = target_s * torch.eye(out_channels * height_out).to( - self.connection.w.device - ) # Update P^+ and P^- values. - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= _cached_decay(self, "tc_plus") self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= _cached_decay(self, "tc_minus") self.p_minus += a_minus * target_s # Calculate point eligibility value. - self.eligibility = torch.bmm(target_s, self.p_plus) + torch.bmm( + self.eligibility = _row_scale(target_s, self.p_plus) + _row_scale( self.p_minus, source_s ) @@ -1728,12 +1876,7 @@ def _local_connection2d_update(self, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = torch.tensor( - kwargs.get("a_plus", 1.0), device=self.connection.w.device - ) - a_minus = torch.tensor( - kwargs.get("a_minus", -1.0), device=self.connection.w.device - ) + a_plus, a_minus = _reward_rates(self, kwargs, self.connection.w.device) # Compute weight update based on the eligibility value of the past timestep. update = reward * self.eligibility @@ -1764,9 +1907,6 @@ def _local_connection2d_update(self, **kwargs) -> None: self.p_minus = self.p_minus.reshape( batch_size, out_channels * height_out * width_out, 1 ) - self.p_minus = self.p_minus * torch.eye( - out_channels * height_out * width_out - ).to(self.connection.w.device) # Reshaping spike occurrences. source_s = ( @@ -1785,18 +1925,15 @@ def _local_connection2d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out * width_out, 1 ) - target_s = target_s * torch.eye(out_channels * height_out * width_out).to( - self.connection.w.device - ) # Update P^+ and P^- values. - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= _cached_decay(self, "tc_plus") self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= _cached_decay(self, "tc_minus") self.p_minus += a_minus * target_s # Calculate point eligibility value. - self.eligibility = torch.bmm(target_s, self.p_plus) + torch.bmm( + self.eligibility = _row_scale(target_s, self.p_plus) + _row_scale( self.p_minus, source_s ) @@ -1831,12 +1968,7 @@ def _local_connection3d_update(self, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = torch.tensor( - kwargs.get("a_plus", 1.0), device=self.connection.w.device - ) - a_minus = torch.tensor( - kwargs.get("a_minus", -1.0), device=self.connection.w.device - ) + a_plus, a_minus = _reward_rates(self, kwargs, self.connection.w.device) # Compute weight update based on the eligibility value of the past timestep. update = reward * self.eligibility @@ -1867,9 +1999,6 @@ def _local_connection3d_update(self, **kwargs) -> None: self.p_minus = self.p_minus.reshape( batch_size, out_channels * height_out * width_out * depth_out, 1 ) - self.p_minus = self.p_minus * torch.eye( - out_channels * height_out * width_out * depth_out - ).to(self.connection.w.device) # Reshaping spike occurrences. source_s = ( @@ -1889,18 +2018,15 @@ def _local_connection3d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out * width_out * depth_out, 1 ) - target_s = target_s * torch.eye( - out_channels * height_out * width_out * depth_out - ).to(self.connection.w.device) # Update P^+ and P^- values. - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= _cached_decay(self, "tc_plus") self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= _cached_decay(self, "tc_minus") self.p_minus += a_minus * target_s # Calculate point eligibility value. - self.eligibility = torch.bmm(target_s, self.p_plus) + torch.bmm( + self.eligibility = _row_scale(target_s, self.p_plus) + _row_scale( self.p_minus, source_s ) @@ -1967,20 +2093,15 @@ def _conv_connection_update(self, dim: int, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = torch.tensor( - kwargs.get("a_plus", 1.0), device=self.connection.w.device - ) - a_minus = torch.tensor( - kwargs.get("a_minus", -1.0), device=self.connection.w.device - ) + a_plus, a_minus = _reward_rates(self, kwargs, self.connection.w.device) source_s = self.source.s.float() target_s = self.target.s.float() def _update_traces_and_eligibility(): - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= _cached_decay(self, "tc_plus") self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= _cached_decay(self, "tc_minus") self.p_minus += a_minus * target_s self.eligibility = _conv_point_eligibility( self.connection, self.p_plus, self.p_minus, source_s, target_s, dim @@ -2061,10 +2182,13 @@ def __init__( self.tc_plus = torch.tensor(kwargs.get("tc_plus", 20.0)) self.tc_minus = torch.tensor(kwargs.get("tc_minus", 20.0)) self.tc_e_trace = torch.tensor(kwargs.get("tc_e_trace", 25.0)) - # If True, the current spikes are folded into the eligibility before the - # eligibility trace is integrated (exact Florian 2007 timing). If False - # (default, backward-compatible) a one-timestep lag is kept. Currently - # honoured by the ``Connection`` path. + # Timing of eligibility vs. trace integration. Default (False) is + # Florian (2007) eqs. 2.7-2.8, z(t+dt) = beta z(t) + zeta(t)/tau_z and + # w(t+dt) = w(t) + gamma dt r(t+dt) z(t+dt): the trace integrated at a + # step uses the eligibility built from the previous step's spikes. + # ``zero_lag=True`` folds this step's spikes into the eligibility + # before integrating (continuous-time eqs. 3.1-3.2 discretised without + # the lag). Currently honoured by the ``Connection`` path. self.zero_lag = kwargs.get("zero_lag", False) def _connection_update(self, **kwargs) -> None: @@ -2106,31 +2230,26 @@ def _connection_update(self, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = torch.tensor( - kwargs.get("a_plus", 1.0), device=self.connection.w.device - ) - a_minus = torch.tensor( - kwargs.get("a_minus", -1.0), device=self.connection.w.device - ) + a_plus, a_minus = _reward_rates(self, kwargs, self.connection.w.device) # Update P^+/P^- traces and the point eligibility for this timestep. def _update_traces_and_eligibility(): - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= _cached_decay(self, "tc_plus") self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= _cached_decay(self, "tc_minus") self.p_minus += a_minus * target_s self.eligibility = torch.bmm( self.p_plus.unsqueeze(2), target_s.unsqueeze(1) ) + torch.bmm(source_s.unsqueeze(2), self.p_minus.unsqueeze(1)) # With zero_lag, fold in the current spikes before integrating the - # eligibility trace (exact Florian 2007 timing). + # eligibility trace (un-lagged variant; the default is eqs. 2.7-2.8). if self.zero_lag: _update_traces_and_eligibility() # Calculate value of eligibility trace based on the value # of the point eligibility value of the past timestep. - self.eligibility_trace *= torch.exp(-self.connection.dt / self.tc_e_trace) + self.eligibility_trace *= _cached_decay(self, "tc_e_trace") self.eligibility_trace += self.eligibility / self.tc_e_trace # Compute weight update, reducing over the minibatch dimension. @@ -2141,8 +2260,8 @@ def _update_traces_and_eligibility(): update = update.to_sparse() self.connection.w += update - # Default (backward-compatible): the eligibility computed here is applied - # on the next timestep (one-step lag). + # Default: the eligibility computed here is applied on the next + # timestep, as in Florian (2007) eq. 3.9 / eqs. 2.7-2.8. if not self.zero_lag: _update_traces_and_eligibility() @@ -2177,16 +2296,11 @@ def _local_connection1d_update(self, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = torch.tensor( - kwargs.get("a_plus", 1.0), device=self.connection.w.device - ) - a_minus = torch.tensor( - kwargs.get("a_minus", -1.0), device=self.connection.w.device - ) + a_plus, a_minus = _reward_rates(self, kwargs, self.connection.w.device) # Integrate the eligibility trace from the point eligibility of the # previous timestep (decay then accumulate). - self.eligibility_trace *= torch.exp(-self.connection.dt / self.tc_e_trace) + self.eligibility_trace *= _cached_decay(self, "tc_e_trace") self.eligibility_trace += self.eligibility / self.tc_e_trace # Compute weight update, reducing over the minibatch dimension. @@ -2212,9 +2326,6 @@ def _local_connection1d_update(self, **kwargs) -> None: self.p_minus = self.p_minus.reshape( batch_size, out_channels * height_out, 1 ) - self.p_minus = self.p_minus * torch.eye(out_channels * height_out).to( - self.connection.w.device - ) # Reshaping spike occurrences. source_s = ( @@ -2229,18 +2340,15 @@ def _local_connection1d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out, 1 ) - target_s = target_s * torch.eye(out_channels * height_out).to( - self.connection.w.device - ) # Update P^+ and P^- values. - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= _cached_decay(self, "tc_plus") self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= _cached_decay(self, "tc_minus") self.p_minus += a_minus * target_s # Calculate point eligibility value. - self.eligibility = torch.bmm(target_s, self.p_plus) + torch.bmm( + self.eligibility = _row_scale(target_s, self.p_plus) + _row_scale( self.p_minus, source_s ) self.eligibility = self.eligibility.view(batch_size, *self.connection.w.shape) @@ -2277,16 +2385,11 @@ def _local_connection2d_update(self, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = torch.tensor( - kwargs.get("a_plus", 1.0), device=self.connection.w.device - ) - a_minus = torch.tensor( - kwargs.get("a_minus", -1.0), device=self.connection.w.device - ) + a_plus, a_minus = _reward_rates(self, kwargs, self.connection.w.device) # Integrate the eligibility trace from the point eligibility of the # previous timestep (decay then accumulate). - self.eligibility_trace *= torch.exp(-self.connection.dt / self.tc_e_trace) + self.eligibility_trace *= _cached_decay(self, "tc_e_trace") self.eligibility_trace += self.eligibility / self.tc_e_trace # Compute weight update, reducing over the minibatch dimension. @@ -2317,9 +2420,6 @@ def _local_connection2d_update(self, **kwargs) -> None: self.p_minus = self.p_minus.reshape( batch_size, out_channels * height_out * width_out, 1 ) - self.p_minus = self.p_minus * torch.eye( - out_channels * height_out * width_out - ).to(self.connection.w.device) # Reshaping spike occurrences. source_s = ( @@ -2339,18 +2439,15 @@ def _local_connection2d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out * width_out, 1 ) - target_s = target_s * torch.eye(out_channels * height_out * width_out).to( - self.connection.w.device - ) # Update P^+ and P^- values. - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= _cached_decay(self, "tc_plus") self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= _cached_decay(self, "tc_minus") self.p_minus += a_minus * target_s # Calculate point eligibility value. - self.eligibility = torch.bmm(target_s, self.p_plus) + torch.bmm( + self.eligibility = _row_scale(target_s, self.p_plus) + _row_scale( self.p_minus, source_s ) self.eligibility = self.eligibility.view(batch_size, *self.connection.w.shape) @@ -2389,16 +2486,11 @@ def _local_connection3d_update(self, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = torch.tensor( - kwargs.get("a_plus", 1.0), device=self.connection.w.device - ) - a_minus = torch.tensor( - kwargs.get("a_minus", -1.0), device=self.connection.w.device - ) + a_plus, a_minus = _reward_rates(self, kwargs, self.connection.w.device) # Integrate the eligibility trace from the point eligibility of the # previous timestep (decay then accumulate). - self.eligibility_trace *= torch.exp(-self.connection.dt / self.tc_e_trace) + self.eligibility_trace *= _cached_decay(self, "tc_e_trace") self.eligibility_trace += self.eligibility / self.tc_e_trace # Compute weight update, reducing over the minibatch dimension. @@ -2430,9 +2522,6 @@ def _local_connection3d_update(self, **kwargs) -> None: self.p_minus = self.p_minus.reshape( batch_size, out_channels * height_out * width_out * depth_out, 1 ) - self.p_minus = self.p_minus * torch.eye( - out_channels * height_out * width_out * depth_out - ).to(self.connection.w.device) # Reshaping spike occurrences. source_s = ( @@ -2453,18 +2542,15 @@ def _local_connection3d_update(self, **kwargs) -> None: target_s = self.target.s.type(torch.float).reshape( batch_size, out_channels * height_out * width_out * depth_out, 1 ) - target_s = target_s * torch.eye( - out_channels * height_out * width_out * depth_out - ).to(self.connection.w.device) # Update P^+ and P^- values. - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= _cached_decay(self, "tc_plus") self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= _cached_decay(self, "tc_minus") self.p_minus += a_minus * target_s # Calculate point eligibility value. - self.eligibility = torch.bmm(target_s, self.p_plus) + torch.bmm( + self.eligibility = _row_scale(target_s, self.p_plus) + _row_scale( self.p_minus, source_s ) self.eligibility = self.eligibility.view(batch_size, *self.connection.w.shape) @@ -2534,20 +2620,15 @@ def _conv_connection_update(self, dim: int, **kwargs) -> None: # Parse keyword arguments. reward = kwargs["reward"] - a_plus = torch.tensor( - kwargs.get("a_plus", 1.0), device=self.connection.w.device - ) - a_minus = torch.tensor( - kwargs.get("a_minus", -1.0), device=self.connection.w.device - ) + a_plus, a_minus = _reward_rates(self, kwargs, self.connection.w.device) source_s = self.source.s.float() target_s = self.target.s.float() def _update_traces_and_eligibility(): - self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus) + self.p_plus *= _cached_decay(self, "tc_plus") self.p_plus += a_plus * source_s - self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus) + self.p_minus *= _cached_decay(self, "tc_minus") self.p_minus += a_minus * target_s self.eligibility = _conv_point_eligibility( self.connection, self.p_plus, self.p_minus, source_s, target_s, dim @@ -2558,7 +2639,7 @@ def _update_traces_and_eligibility(): _update_traces_and_eligibility() # Integrate the eligibility trace and apply the weight update. - self.eligibility_trace *= torch.exp(-self.connection.dt / self.tc_e_trace) + self.eligibility_trace *= _cached_decay(self, "tc_e_trace") self.eligibility_trace += self.eligibility / self.tc_e_trace self.connection.w += ( self.nu[0] * self.connection.dt * reward * self.eligibility_trace @@ -2601,8 +2682,9 @@ def __init__( Keyword arguments: - :param float tc_c: Time constant for balancing naive Hebbian and policy gradient - learning. + :param float tc_c: Time constant :math:`\\tau_c` balancing policy-gradient and + naive Hebbian learning (Vasilaki et al. 2009, eq. 8): ``0`` gives the + strict policy-gradient rule, ``inf`` the naive Hebbian rule. Default 5. :param float tc_e_trace: Time constant for the eligibility trace. """ super().__init__( @@ -2632,7 +2714,7 @@ def __init__( self.tc_c = torch.tensor( kwargs.get("tc_c", 5.0) - ) # 0 for pure naive Hebbian, inf for pure policy gradient. + ) # 0 for strict policy gradient, inf for pure naive Hebbian (eq. 8). self.tc_e_trace = torch.tensor(kwargs.get("tc_e_trace", 25.0)) def _connection_update(self, **kwargs) -> None: diff --git a/bindsnet/models/models.py b/bindsnet/models/models.py index 04ff33612..a5027e0c6 100644 --- a/bindsnet/models/models.py +++ b/bindsnet/models/models.py @@ -7,6 +7,7 @@ from torch import device from bindsnet.learning import PostPre +from bindsnet.learning.MCC_learning import DiehlAndCook as MMCDiehlAndCook from bindsnet.learning.MCC_learning import PostPre as MMCPostPre from bindsnet.network import Network from bindsnet.network.nodes import DiehlAndCookNodes, Input, LIFNodes @@ -119,6 +120,8 @@ def __init__( inpt_shape: Optional[Iterable[int]] = None, inh_thresh: float = -40.0, exc_thresh: float = -52.0, + learning_rule: Optional[type] = None, + learning_rule_kwargs: Optional[dict] = None, ) -> None: # language=rst """ @@ -143,6 +146,14 @@ def __init__( :param tc_theta_decay: Time constant of ``DiehlAndCookNodes`` threshold potential decay. :param inpt_shape: The dimensionality of the input layer. + :param learning_rule: ``bindsnet.learning.MCC_learning`` rule for the input to + excitatory weights. Default ``PostPre`` (pair-based STDP, the rule the + published BindsNET replication used). Pass + ``bindsnet.learning.MCC_learning.DiehlAndCook`` for the paper's own + post-spike-only rule; that also switches the input traces to the paper's + accumulating form (``traces_additive=True``). + :param learning_rule_kwargs: Extra options for the rule, e.g. + ``{"x_tar": 0.4, "mu": 1.0}`` for ``DiehlAndCook``. """ super().__init__(dt=dt) @@ -153,9 +164,18 @@ def __init__( self.inh = inh self.dt = dt + if learning_rule is None: + learning_rule = MMCPostPre + learning_rule_kwargs = dict(learning_rule_kwargs or {}) + paper_rule = learning_rule is MMCDiehlAndCook + # Layers input_layer = Input( - n=self.n_inpt, shape=self.inpt_shape, traces=True, tc_trace=20.0 + n=self.n_inpt, + shape=self.inpt_shape, + traces=True, + traces_additive=paper_rule, + tc_trace=20.0, ) exc_layer = DiehlAndCookNodes( n=self.n_neurons, @@ -195,9 +215,10 @@ def __init__( norm=norm, reduction=reduction, nu=nu, - learning_rule=MMCPostPre, + learning_rule=learning_rule, sparse=sparse, batch_size=batch_size, + **learning_rule_kwargs, ) ], ) diff --git a/bindsnet/network/network.py b/bindsnet/network/network.py index 0359e9c1f..7754bf890 100644 --- a/bindsnet/network/network.py +++ b/bindsnet/network/network.py @@ -9,18 +9,46 @@ from bindsnet.network.topology import AbstractConnection -def load(file_name: str, map_location: str = "cpu", learning: bool = None) -> "Network": +def load( + file_name: str, + map_location: str = "cpu", + learning: bool = None, + weights_only: bool = False, +) -> "Network": # language=rst """ Loads serialized network object from disk. + .. warning:: + **Only load network files you created yourself or otherwise trust.** + + Network files are Python pickle files. Loading one runs code that is + stored inside it, so a file from an untrusted source (a download, a + model-sharing site, shared group storage, a file sent by someone else) + can run arbitrary commands on your machine as soon as you call this + function. There is no way to inspect a pickle file safely before + loading it. + + This is the standard behaviour of :func:`torch.load` and applies to + saved models across the PyTorch ecosystem, not just to BindsNET. + + ``weights_only=True`` asks PyTorch to refuse to execute code while + loading, but it **cannot** load networks written by + :py:meth:`Network.save`, because those store the whole network object + rather than a plain tensor state dictionary. It is offered here only + for files you know contain plain tensors. + :param file_name: Path to serialized network object on disk. :param map_location: One of ``"cpu"`` or ``"cuda"``. Defaults to ``"cpu"``. :param learning: Whether to load with learning enabled. Default loads value from disk. + :param weights_only: Passed through to :func:`torch.load`. Defaults to + ``False``, which permits code execution and is required to load files + written by :py:meth:`Network.save`. Set to ``True`` to refuse code + execution, which only works for files holding plain tensors. """ network = torch.load( - open(file_name, "rb"), map_location=map_location, weights_only=False + open(file_name, "rb"), map_location=map_location, weights_only=weights_only ) if learning is not None and "learning" in vars(network): network.learning = learning @@ -165,6 +193,12 @@ def save(self, file_name: str) -> None: """ Serializes the network object to disk. + .. note:: + The saved file is a Python pickle file holding the whole network + object. Anyone who loads it with :py:func:`bindsnet.network.load` + runs whatever code it contains, so treat a network file you share + the same way you would treat a script you ask someone to run. + :param file_name: Path to store serialized network object on disk. **Example:** @@ -193,7 +227,6 @@ def save(self, file_name: str) -> None: # Save the network to disk. network.save(str(Path.home()) + '/network.pt') """ - torch.serialization.add_safe_globals([self]) torch.save(self, open(file_name, "wb")) def clone(self) -> "Network": @@ -206,7 +239,11 @@ def clone(self) -> "Network": virtual_file = tempfile.SpooledTemporaryFile() torch.save(self, virtual_file) virtual_file.seek(0) - return torch.load(virtual_file) + # weights_only=False is required: this buffer holds the whole network + # object, written by torch.save just above, so it is trusted by + # construction. PyTorch 2.6 changed the default to True, which cannot + # read it. + return torch.load(virtual_file, weights_only=False) def _get_inputs(self, layers: Iterable = None) -> Dict[str, torch.Tensor]: # language=rst @@ -267,10 +304,15 @@ def run( :param Dict[str, torch.Tensor] clamp: Mapping of layer names to boolean masks if neurons should be clamped to spiking. The ``Tensor``s have shape - ``[n_neurons]`` or ``[time, n_neurons]``. + ``[n_neurons]`` or ``[time, n_neurons]``. A clamped spike is a real spike + for everything downstream: it is propagated through connections, enters + the layer's spike trace (``x``) and therefore the learning rules, and is + recorded by monitors. It does not reset the neuron's voltage or start its + refractory period. :param Dict[str, torch.Tensor] unclamp: Mapping of layer names to boolean masks if neurons should be clamped to not spiking. The ``Tensor``s should have - shape ``[n_neurons]`` or ``[time, n_neurons]``. + shape ``[n_neurons]`` or ``[time, n_neurons]``. A suppressed spike is + removed before the trace update, so it leaves no trace. :param Dict[str, torch.Tensor] injects_v: Mapping of layer names to boolean masks if neurons should be added voltage. The ``Tensor``s should have shape ``[n_neurons]`` or ``[time, n_neurons]``. @@ -403,6 +445,21 @@ def run( else: self.layers[l].v += inject_v[t] + # Spike clamps for this step. The layer applies them inside + # ``forward`` before updating its spike traces, so forced / + # suppressed spikes are seen by the traces and hence by the + # learning rules (see ``Nodes.forward``). + clamp = clamps.get(l, None) + if clamp is not None: + self.layers[l]._clamp = ( + clamp if clamp.ndimension() == 1 else clamp[t] + ) + unclamp = unclamps.get(l, None) + if unclamp is not None: + self.layers[l]._unclamp = ( + unclamp if unclamp.ndimension() == 1 else unclamp[t] + ) + if l in current_inputs: self.layers[l].forward(x=current_inputs[l]) else: @@ -412,22 +469,6 @@ def run( ) ) - # Clamp neurons to spike. - clamp = clamps.get(l, None) - if clamp is not None: - if clamp.ndimension() == 1: - self.layers[l].s[:, clamp] = 1 - else: - self.layers[l].s[:, clamp[t]] = 1 - - # Clamp neurons not to spike. - unclamp = unclamps.get(l, None) - if unclamp is not None: - if unclamp.ndimension() == 1: - self.layers[l].s[:, unclamp] = 0 - else: - self.layers[l].s[:, unclamp[t]] = 0 - for c in self.connections: flad_m = False if A_Minus != None and ((isinstance(A_Minus, float)) or (c in A_Minus)): diff --git a/bindsnet/network/nodes.py b/bindsnet/network/nodes.py index cf8b709cb..1fcdd3af9 100644 --- a/bindsnet/network/nodes.py +++ b/bindsnet/network/nodes.py @@ -12,6 +12,13 @@ class Nodes(torch.nn.Module): Abstract base class for groups of neurons. """ + # Per-step spike clamps set by ``Network.run`` (``clamp`` / ``unclamp`` + # keyword arguments). They are applied in ``forward`` *before* the spike + # traces are updated, so a forced or suppressed spike is reflected in the + # traces that the learning rules read. Cleared after use. + _clamp: Optional[torch.Tensor] = None + _unclamp: Optional[torch.Tensor] = None + def __init__( self, n: Optional[int] = None, @@ -93,18 +100,27 @@ def forward(self, x: torch.Tensor) -> None: :param x: Inputs to the layer. """ + # Force / suppress spikes requested for this step (see ``Network.run``), + # before the traces below see ``self.s``. + if self._clamp is not None: + self.s[:, self._clamp] = 1 + self._clamp = None + if self._unclamp is not None: + self.s[:, self._unclamp] = 0 + self._unclamp = None + if self.traces: # Decay and set spike traces. - self.x *= self.trace_decay + self.x.mul_(self.trace_decay) if self.traces_additive: - self.x += self.trace_scale * self.s.float() + self.x.add_(self.trace_scale * self.s.float()) else: self.x.masked_fill_(self.s.bool(), self.trace_scale) if self.sum_input: # Add current input to running sum. - self.summed += x.float() + self.summed.add_(x.float()) def reset_state_variables(self) -> None: # language=rst @@ -376,10 +392,10 @@ def forward(self, x: torch.Tensor) -> None: :param x: Inputs to the layer. """ # Integrate input voltages. - self.v += (self.refrac_count <= 0).float() * x + self.v.add_((self.refrac_count <= 0).float() * x) # Decrement refractory counters. - self.refrac_count -= self.dt + self.refrac_count.sub_(self.dt) # Check for spiking neurons. self.s = self.v >= self.thresh @@ -505,15 +521,17 @@ def forward(self, x: torch.Tensor) -> None: :param x: Inputs to the layer. """ # Decay voltages. - self.v = self.decay * (self.v - self.rest) + self.rest + # In place: same three operations in the same order as + # ``decay * (v - rest) + rest`` (bit-identical), without the temporaries. + self.v.sub_(self.rest).mul_(self.decay).add_(self.rest) # Integrate inputs. x.masked_fill_(self.refrac_count > 0, 0.0) # Decrement refractory counters. - self.refrac_count -= self.dt + self.refrac_count.sub_(self.dt) - self.v += x # interlaced + self.v.add_(x) # interlaced # Check for spiking neurons. self.s = self.v >= self.thresh @@ -626,17 +644,17 @@ def forward(self, x: torch.Tensor) -> None: :param x: Inputs to the layer. """ # Decay voltages. - self.v *= self.decay + self.v.mul_(self.decay) # Integrate inputs. if x is not None: x.masked_fill_(self.refrac_count > 0, 0.0) # Decrement refractory counters. - self.refrac_count -= self.dt + self.refrac_count.sub_(self.dt) if x is not None: - self.v += x + self.v.add_(x) # Check for spiking neurons. self.s = self.v >= self.thresh @@ -767,15 +785,17 @@ def forward(self, x: torch.Tensor) -> None: :param x: Inputs to the layer. """ # Decay voltages and current. - self.v = self.decay * (self.v - self.rest) + self.rest - self.i *= self.i_decay + # In place: same three operations in the same order as + # ``decay * (v - rest) + rest`` (bit-identical), without the temporaries. + self.v.sub_(self.rest).mul_(self.decay).add_(self.rest) + self.i.mul_(self.i_decay) # Decrement refractory counters. - self.refrac_count -= self.dt + self.refrac_count.sub_(self.dt) # Integrate inputs. - self.i += x - self.v += (self.refrac_count <= 0).float() * self.i + self.i.add_(x) + self.v.add_((self.refrac_count <= 0).float() * self.i) # Check for spiking neurons. self.s = self.v >= self.thresh @@ -919,15 +939,17 @@ def forward(self, x: torch.Tensor) -> None: :param x: Inputs to the layer. """ # Decay voltages and adaptive thresholds. - self.v = self.decay * (self.v - self.rest) + self.rest + # In place: same three operations in the same order as + # ``decay * (v - rest) + rest`` (bit-identical), without the temporaries. + self.v.sub_(self.rest).mul_(self.decay).add_(self.rest) if self.learning: - self.theta *= self.theta_decay + self.theta.mul_(self.theta_decay) # Integrate inputs. - self.v += (self.refrac_count <= 0).float() * x + self.v.add_((self.refrac_count <= 0).float() * x) # Decrement refractory counters. - self.refrac_count -= self.dt + self.refrac_count.sub_(self.dt) # Check for spiking neurons. self.s = self.v >= self.thresh + self.theta @@ -936,7 +958,7 @@ def forward(self, x: torch.Tensor) -> None: self.refrac_count.masked_fill_(self.s, self.refrac) self.v.masked_fill_(self.s, self.reset) if self.learning: - self.theta += self.theta_plus * self.s.float().sum(0) + self.theta.add_(self.theta_plus * self.s.float().sum(0)) # voltage clipping to lowerbound if self.lbound is not None: @@ -1074,15 +1096,17 @@ def forward(self, x: torch.Tensor) -> None: :param x: Inputs to the layer. """ # Decay voltages and adaptive thresholds. - self.v = self.decay * (self.v - self.rest) + self.rest + # In place: same three operations in the same order as + # ``decay * (v - rest) + rest`` (bit-identical), without the temporaries. + self.v.sub_(self.rest).mul_(self.decay).add_(self.rest) if self.learning: - self.theta *= self.theta_decay + self.theta.mul_(self.theta_decay) # Integrate inputs. - self.v += (self.refrac_count <= 0).float() * x + self.v.add_((self.refrac_count <= 0).float() * x) # Decrement refractory counters. - self.refrac_count -= self.dt + self.refrac_count.sub_(self.dt) # Check for spiking neurons. self.s = self.v >= self.thresh + self.theta @@ -1091,7 +1115,7 @@ def forward(self, x: torch.Tensor) -> None: self.refrac_count.masked_fill_(self.s, self.refrac) self.v.masked_fill_(self.s, self.reset) if self.learning: - self.theta += self.theta_plus * self.s.float().sum(0) + self.theta.add_(self.theta_plus * self.s.float().sum(0)) # Choose only a single neuron to spike. if self.one_spike: @@ -1282,9 +1306,9 @@ def forward(self, x: torch.Tensor) -> None: ) # Apply v and u updates. - self.v += self.dt * 0.5 * (0.04 * self.v**2 + 5 * self.v + 140 - self.u + x) - self.v += self.dt * 0.5 * (0.04 * self.v**2 + 5 * self.v + 140 - self.u + x) - self.u += self.dt * self.a * (self.b * self.v - self.u) + self.v.add_(self.dt * 0.5 * (0.04 * self.v**2 + 5 * self.v + 140 - self.u + x)) + self.v.add_(self.dt * 0.5 * (0.04 * self.v**2 + 5 * self.v + 140 - self.u + x)) + self.u.add_(self.dt * self.a * (self.b * self.v - self.u)) # Voltage clipping to lower bound. if self.lbound is not None: @@ -1432,10 +1456,10 @@ def forward(self, x: torch.Tensor) -> None: :param x: Inputs to the layer. """ # Decay voltages. - self.v *= self.decay + self.v.mul_(self.decay) if self.learning: - self.theta *= self.theta_decay + self.theta.mul_(self.theta_decay) # Integrate inputs. v = torch.einsum( @@ -1444,13 +1468,13 @@ def forward(self, x: torch.Tensor) -> None: v += torch.einsum( "i,kij->kj", self.refKernel, self.last_spikes ) # Refractoriness due to previous spikes - self.v += v.view(x.size(0), *self.shape) + self.v.add_(v.view(x.size(0), *self.shape)) # Check for spiking neurons. self.s = self.v >= self.thresh + self.theta if self.learning: - self.theta += self.theta_plus * self.s.float().sum(0) + self.theta.add_(self.theta_plus * self.s.float().sum(0)) # Add the spike vector into the first in first out matrix of windowed (ref) spike trains self.last_spikes = torch.cat( @@ -1644,10 +1668,12 @@ def forward(self, x: torch.Tensor) -> None: :param x: Inputs to the layer. """ # Decay voltages. - self.v = self.decay * (self.v - self.rest) + self.rest + # In place: same three operations in the same order as + # ``decay * (v - rest) + rest`` (bit-identical), without the temporaries. + self.v.sub_(self.rest).mul_(self.decay).add_(self.rest) # Integrate inputs. - self.v += (self.refrac_count <= 0).float() * self.eps_0 * x + self.v.add_((self.refrac_count <= 0).float() * self.eps_0 * x) # Compute (instantaneous) probabilities of spiking, clamp between 0 and 1 using exponentials. # Also known as 'escape noise', this simulates nearby neurons. @@ -1655,7 +1681,7 @@ def forward(self, x: torch.Tensor) -> None: self.s_prob = 1.0 - torch.exp(-self.rho * self.dt) # Decrement refractory counters. - self.refrac_count -= self.dt + self.refrac_count.sub_(self.dt) # Check for spiking neurons (spike when probability > some random number). self.s = torch.rand_like(self.s_prob) < self.s_prob diff --git a/bindsnet/network/topology.py b/bindsnet/network/topology.py index f277e831b..9c342a6c3 100644 --- a/bindsnet/network/topology.py +++ b/bindsnet/network/topology.py @@ -261,7 +261,7 @@ def remove_pipeline(self, feature) -> None: self.pipeline.remove(feature) del self.feature_index[feature.name] - def _apply(self, fn): + def _apply(self, fn, recurse=True): # language=rst """ Relocate pipeline features (and their learning rules) along with the connection @@ -274,7 +274,7 @@ def _apply(self, fn): mismatch. The feature value is moved in place (via ``.data``) so it stays aliased to the learning rule's cached reference. """ - super()._apply(fn) + super()._apply(fn, recurse) for feature in self.pipeline: value = getattr(feature, "value", None) if isinstance(value, torch.Tensor): diff --git a/bindsnet/network/topology_features.py b/bindsnet/network/topology_features.py index 76c885899..aa6c16f8d 100644 --- a/bindsnet/network/topology_features.py +++ b/bindsnet/network/topology_features.py @@ -92,6 +92,7 @@ def __init__( NoOp, PostPre, Hebbian, + DiehlAndCook, MSTDP, MSTDPET, ) @@ -100,6 +101,7 @@ def __init__( NoOp, PostPre, Hebbian, + DiehlAndCook, MSTDP, MSTDPET, ] @@ -246,6 +248,8 @@ def prime_feature(self, connection, device, **kwargs) -> None: if self.learning_rule is None: self.learning_rule = NoOp + # Rule-specific options given to the feature (e.g. ``x_tar``, ``mu``, + # ``tc_plus``) are forwarded to the rule; connection-level kwargs win. self.learning_rule = self.learning_rule( connection=connection, feature_value=self.value, @@ -253,7 +257,7 @@ def prime_feature(self, connection, device, **kwargs) -> None: nu=self.nu, reduction=self.reduction, decay=self.decay, - **kwargs, + **{**self.kwargs, **kwargs}, ) #### Recycle unnecessary variables #### @@ -595,6 +599,7 @@ def __init__( decay: float = 0.0, sparse: Optional[bool] = False, batch_size: int = 1, + **kwargs, ) -> None: # language=rst """ @@ -616,6 +621,9 @@ def __init__( :param decay: Constant multiple to decay weights by on each iteration :param sparse: Should :code:`value` parameter be sparse tensor or not :param batch_size: Mini-batch size. + + Any further keyword arguments (e.g. ``x_tar``, ``mu``, ``tc_plus``) are + forwarded to the learning rule when it is instantiated. """ self.norm_frequency = norm_frequency @@ -638,6 +646,7 @@ def __init__( decay=decay, sparse=sparse, batch_size=batch_size, + **kwargs, ) def compute(self, s) -> Union[torch.Tensor, float, int]: diff --git a/docs/source/installation.rst b/docs/source/installation.rst index e1355d590..f0cd208f0 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -5,7 +5,13 @@ Installation Pip install ----------- -Issue: +To install the latest release from `PyPI `_, issue: + +.. code-block:: bash + + pip install bindsnet + +To install the current development code from GitHub, issue: .. code-block:: bash @@ -19,8 +25,8 @@ On \*nix systems, issue one of the following in a shell: .. code-block:: bash - git clone https://github.com/Hananel-Hazan/bindsnet.git # HTTPS - git clone git@github.com:Hananel-Hazan/bindsnet.git # SSH + git clone https://github.com/BindsNET/bindsnet.git # HTTPS + git clone git@github.com:BindsNET/bindsnet.git # SSH Change directory into :code:`bindsnet` and issue one of the following: diff --git a/docs/source/models_spec.rst b/docs/source/models_spec.rst index a4e64118d..771de9814 100644 --- a/docs/source/models_spec.rst +++ b/docs/source/models_spec.rst @@ -177,6 +177,20 @@ i.e. a pre-synaptic spike **depresses** the synapse in proportion to the post-sy trace, and a post-synaptic spike **potentiates** it in proportion to the pre-synaptic trace. Convolutional and locally-connected variants apply the same rule patch-wise. +This is the additive pair-based trace STDP of Morrison, Diesmann & Gerstner (2008), +*Biol. Cybern.* 98:459-478, eqs. (11)-(14), with :math:`F_+ = \nu_\text{post}`, +:math:`F_- = \nu_\text{pre}`. Traces follow their Sect. 2.3: ``traces_additive=True`` +accumulates 1 per spike; ``traces_additive=False`` resets the trace to 1 on each spike. +Validated in ``test/network/test_learning_rule_specs.py``. + +.. note:: + + ``PostPre`` is **not** the rule of Diehl & Cook (2015) even though it is the rule + used by the ``DiehlAndCook2015`` model. Diehl & Cook change weights only on + post-synaptic spikes, :math:`\Delta w = \eta (x_\text{pre} - x_\text{tar})(w_\max - w)^\mu`, + with a target trace :math:`x_\text{tar}`; ``PostPre`` has no target trace and adds a + depression term on pre-synaptic spikes instead. + Hebbian (``Hebbian``) ~~~~~~~~~~~~~~~~~~~~~~ Both pre- and post-synaptic events **increase** the weight (no depression term), @@ -185,7 +199,29 @@ proportional to the opposite layer's trace. Weight-dependent post-pre (``WeightDependentPostPre``) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ``PostPre`` whose potentiation/depression magnitudes are scaled by the distance of the -weight from its bounds (``wmin``/``wmax``), yielding soft saturation at the limits. +weight from its bounds (``wmin``/``wmax``), yielding soft saturation at the limits: +Morrison et al. (2008) eqs. (13)-(14) with :math:`F_+ = \nu_\text{post}(w_\max - w)` and +:math:`F_- = \nu_\text{pre}(w - w_\min)` (the multiplicative / soft-bound rule). + +Diehl & Cook STDP (``DiehlAndCook``) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +The post-spike-only rule of Diehl & Cook (2015), *Front. Comput. Neurosci.* 9:99, +Sect. 2.3 "Learning": on every post-synaptic spike + +.. math:: + + \\Delta w = \\eta\\,(x_\\text{pre} - x_\\text{tar})\\,(w_\\max - w)^\\mu + +where the pre-synaptic trace :math:`x_\\text{pre}` is increased by 1 on each +pre-synaptic spike and decays exponentially (``traces_additive=True``), +:math:`x_\\text{tar}` is the target trace value ("the higher the target value, the lower +the synaptic weight will be"), :math:`w_\\max` is ``wmax`` and :math:`\\mu` the weight +dependence. Pre-synaptic spikes do not change the weight. Keyword arguments ``x_tar`` +(default 0) and ``mu`` (default 1); the paper gives no numeric values for them. Only +the post-synaptic learning rate ``nu[1]`` is used. Available for ``Connection`` / +``LocalConnection`` and as ``MCC_learning.DiehlAndCook`` for multicompartment +connections; ``DiehlAndCook2015(learning_rule=MCC_learning.DiehlAndCook, ...)`` +switches the model to it. Validated in ``test/network/test_learning_rule_specs.py``. Reward-modulated STDP (``MSTDP``, ``MSTDPET``) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -193,16 +229,53 @@ Three-factor rules: a STDP-like eligibility signal is gated by a scalar **reward ``MSTDP`` modulates the immediate pre/post correlation by reward; ``MSTDPET`` adds an **eligibility trace** that accumulates the correlation over time (time constant ``tc_e_trace``) before reward gating. Reward is supplied via the pipeline / an -``AbstractReward`` (e.g. ``MovingAvgRPE``). See source for the exact eligibility update. +``AbstractReward`` (e.g. ``MovingAvgRPE``). + +Both follow the discrete-time equations of Florian (2007), *Neural Comput.* +19:1468-1502: traces (3.11)-(3.12), eligibility (3.10), ``MSTDP`` update (3.9) +:math:`w(t+\delta t) = w(t) + \gamma\, r(t+\delta t)\, \zeta(t)`, and ``MSTDPET`` +(2.7)-(2.8). The reward passed to ``network.run`` at a step therefore multiplies the +eligibility of the previous step (``zero_lag=False``, the default). Defaults +``tc_plus = tc_minus = 20``, ``tc_e_trace = 25``, ``a_plus = 1``, ``a_minus = -1`` are +the paper's. Validated in ``test/network/test_mstdp_florian.py``. Rmax (``Rmax``) ~~~~~~~~~~~~~~~ -Reward-maximizing rule intended for stochastic (SRM0) neurons; see source for its -formulation. +Reward-maximizing rule for stochastic ``SRM0Nodes``: Vasilaki, Fremaux, Urbanczik, +Senn & Gerstner (2009), *PLoS Comput. Biol.* 5(12):e1000586, eqs. (7)-(8) with the +escape rate of eq. (13). The eligibility trace decays with ``tc_e_trace`` and, on each +step, adds :math:`[Y_i - p_i / (1 + (\tau_c/\delta t)\, p_i)]` times the additive +pre-synaptic trace, where :math:`Y_i` is the post-synaptic spike and +:math:`p_i = 1 - e^{-\rho_i \delta t}` its spike probability; ``tc_c`` is +:math:`\tau_c` (``0`` = strict policy gradient, ``inf`` = naive Hebbian). The constant +:math:`g'/g = 1/\Delta u` of eq. (8) is absorbed into ``nu``. Validated in +``test/network/test_learning_rule_specs.py``. + +Spike clamps and traces +~~~~~~~~~~~~~~~~~~~~~~~ +``network.run(..., clamp={layer: mask}, unclamp={layer: mask})`` forces or suppresses +spikes for a step. The clamp is applied inside ``Nodes.forward`` *before* the spike +trace ``x`` is updated, so a forced spike leaves a trace and a suppressed spike does +not; the learning rules therefore treat clamped spikes exactly like natural ones. A +forced spike does not reset the neuron's voltage or start its refractory period. +(Before September 2026 the clamp was applied after the trace update, so clamped +spikes entered the same-step potentiation term but never the trace.) + +Known deviations from the papers +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +* ``PostPre`` is not the Diehl & Cook (2015) rule (see the note above). The + ``DiehlAndCook2015`` model keeps ``PostPre`` as its default because that is what the + published BindsNET replication used; the paper's rule is available as + ``DiehlAndCook`` and can be selected with the model's ``learning_rule`` argument. +* Until September 2026 the multicompartment ``MCC_learning.PostPre`` multiplied each + update by the simulation step ``dt`` (the classic ``PostPre`` and Morrison et al. + (2008) eqs. (13)-(14) do not). The factor was removed; results at ``dt = 1`` are + unchanged, and at other steps the update is now a per-spike increment like the + classic rule. .. note:: - Where this page summarizes a rule "see source", the equations were not reproduced here - to avoid mis-stating constants; consult ``bindsnet/learning/learning.py`` for the - authoritative form. If an implementation deviates from a textbook model, the code is - the specification. + Where this page does not reproduce an equation, consult + ``bindsnet/learning/learning.py`` for the authoritative form. If an implementation + deviates from a textbook model, the code is the specification, and the deviation + is listed above. diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst index 117c542e9..286fa8071 100644 --- a/docs/source/quickstart.rst +++ b/docs/source/quickstart.rst @@ -2,7 +2,7 @@ Quickstart ========== Check out some example use cases for BindsNET in the :code:`examples/` folder -(`link `_). For example, changing directory to +(`link `_). For example, changing directory to `[bindsnet-root]/examples/mnist` and running the following will result in a near-replication of the architecture of `Diehl & Cook 2015 `_: diff --git a/examples/benchmark/hot_path_bench.py b/examples/benchmark/hot_path_bench.py new file mode 100644 index 000000000..27c7f053c --- /dev/null +++ b/examples/benchmark/hot_path_bench.py @@ -0,0 +1,120 @@ +""" +Timing / peak-memory benchmark for the per-timestep hot paths (learning rules, +neuron updates). Prints one line per workload. + + python examples/benchmark/hot_path_bench.py # CPU + python examples/benchmark/hot_path_bench.py --gpu # CUDA +""" + +import argparse +import time + +import torch + +from bindsnet.encoding import poisson +from bindsnet.learning import MSTDP, PostPre +from bindsnet.models import DiehlAndCook2015 +from bindsnet.network import Network +from bindsnet.network.nodes import Input, LIFNodes +from bindsnet.network.topology import Connection, LocalConnection2D + +parser = argparse.ArgumentParser() +parser.add_argument("--gpu", action="store_true") +parser.add_argument("--reps", type=int, default=3) +parser.add_argument("--time", type=int, default=250) +args = parser.parse_args() +dev = "cuda" if args.gpu and torch.cuda.is_available() else "cpu" +T = args.time + + +def timeit(net, inputs, **kw): + if dev != "cpu": + net.to(dev) + inputs = {k: v.to(dev) for k, v in inputs.items()} + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + net.run(inputs=inputs, time=T, **kw) # warm-up + net.reset_state_variables() + if dev != "cpu": + torch.cuda.synchronize() + t = time.perf_counter() + for _ in range(args.reps): + net.run(inputs=inputs, time=T, **kw) + net.reset_state_variables() + if dev != "cpu": + torch.cuda.synchronize() + wall = (time.perf_counter() - t) / args.reps + mem = torch.cuda.max_memory_allocated() / 2**20 if dev != "cpu" else float("nan") + return wall, mem + + +rows = [] + +for b in (1, 16): + torch.manual_seed(0) + net = DiehlAndCook2015( + n_inpt=784, n_neurons=400, batch_size=b, inpt_shape=(1, 28, 28), device=dev + ) + inp = poisson(torch.rand(b, 1, 28, 28) * 128, time=T) + rows.append((f"DiehlAndCook2015 784->400 batch={b}",) + timeit(net, {"X": inp})) + +for b in (1, 16): + torch.manual_seed(0) + net = Network(dt=1.0, batch_size=b) + net.add_layer(Input(n=784, traces=True), "in") + net.add_layer(LIFNodes(n=1000, traces=True), "out") + net.add_connection( + Connection( + net.layers["in"], + net.layers["out"], + nu=(1e-4, 1e-2), + update_rule=PostPre, + wmin=0, + wmax=1, + norm=78.4, + ), + "in", + "out", + ) + inp = poisson(torch.rand(b, 784) * 128, time=T) + rows.append((f"Connection+PostPre 784->1000 batch={b}",) + timeit(net, {"in": inp})) + +torch.manual_seed(0) +net = Network(dt=1.0) +net.add_layer(Input(shape=[1, 28, 28], traces=True), "in") +net.add_layer(LIFNodes(shape=[16, 6, 6], traces=True), "out") +net.add_connection( + LocalConnection2D( + net.layers["in"], + net.layers["out"], + kernel_size=8, + stride=4, + n_filters=16, + nu=(1e-4, 1e-2), + update_rule=PostPre, + wmin=0, + wmax=1, + ), + "in", + "out", +) +inp = poisson(torch.rand(1, 1, 28, 28) * 128, time=T) +rows.append(("LocalConnection2D+PostPre 28x28 k8 s4 f16",) + timeit(net, {"in": inp})) + +torch.manual_seed(0) +net = Network(dt=1.0) +net.add_layer(Input(n=784), "in") +net.add_layer(LIFNodes(n=500), "out") +net.add_connection( + Connection( + net.layers["in"], net.layers["out"], nu=1e-3, update_rule=MSTDP, wmin=0, wmax=1 + ), + "in", + "out", +) +inp = poisson(torch.rand(1, 784) * 128, time=T) +rows.append(("Connection+MSTDP 784->500",) + timeit(net, {"in": inp}, reward=1.0)) + +print(f"device={dev} timesteps={T} reps={args.reps}") +for name, wall, mem in rows: + print(f"{name:<45s} {wall:8.3f} s/run peak {mem:8.1f} MiB") diff --git a/poetry.lock b/poetry.lock index 4d8797b87..01940749e 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.4.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.4.3 and should not be changed by hand. [[package]] name = "ale-py" @@ -329,7 +329,7 @@ version = "2026.7.22" description = "Python package for providing Mozilla's CA Bundle." optional = false python-versions = ">=3.7" -groups = ["main", "dev"] +groups = ["dev"] files = [ {file = "certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775"}, {file = "certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55"}, @@ -466,7 +466,7 @@ version = "3.5.1" description = "The Real First Universal Charset Detector. Open, modern and actively maintained alternative to Chardet." optional = false python-versions = ">=3.7" -groups = ["main", "dev"] +groups = ["dev"] files = [ {file = "charset_normalizer-3.5.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:d1ee1e296209fdce05b81b663250eefa02213a2da7b41bf26f7829b8ba3545aa"}, {file = "charset_normalizer-3.5.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e9fbdce1e47394b09bc9f26ab117dfc8d6491977a11d86f592bb42c779db2fda"}, @@ -902,55 +902,6 @@ files = [ docs = ["ipython", "matplotlib", "numpydoc", "sphinx"] tests = ["pytest", "pytest-cov", "pytest-xdist"] -[[package]] -name = "cython" -version = "3.3.0" -description = "The Cython compiler for writing C extensions in the Python language." -optional = false -python-versions = ">=3.9" -groups = ["main"] -files = [ - {file = "cython-3.3.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0507d9caf7dc35f1212627145d5d13dbc5dd7128529a6608ab72690472fa688e"}, - {file = "cython-3.3.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:de883ec6764b61547c1e7674c0d8a8a875d398bd6bb684e46b93d83e4f13b260"}, - {file = "cython-3.3.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:eda47eb7731c3b41180b58bb83de423f43aa58a677677e3390e8d332b003859e"}, - {file = "cython-3.3.0-cp310-cp310-win_amd64.whl", hash = "sha256:bf411da3ef1af8763781c219108860f7de33f1100038da35d6bf1b4d83fcb2c0"}, - {file = "cython-3.3.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:ec09dbf73ff4f7be2b339b995fadae9c4bb517bbbed7ec11d6fe99c2092b48fd"}, - {file = "cython-3.3.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:11e437f086affee8051cec4bb531be3edb646ab66e325154aa6849377f365033"}, - {file = "cython-3.3.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e6035b5231a9316edc19d6415f4296fd1d0370e2a165a714b3edc167b9ca00e1"}, - {file = "cython-3.3.0-cp311-cp311-win_amd64.whl", hash = "sha256:8566ea804cfc265f5e9dda71d1b716aa24ee4c3423a5da4b28a248a78c33e3f9"}, - {file = "cython-3.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:03bc5333932f5dda3ba9315298ecdd21daa1b58410bb1f8ce04c78ec8337130a"}, - {file = "cython-3.3.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e321ae700995a16dc3055ada06ffb8d61e1a7434e5d0e811547a45ac1015ebd"}, - {file = "cython-3.3.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:428fafed98ea26927000a287b4dfc9ef07339f56656a5329a34eaa593f79a4f8"}, - {file = "cython-3.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:333449cc0350baedee5a6af27929eac8a71eac4ec59333c45ff476b33c6c660d"}, - {file = "cython-3.3.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:03056533fe4fdbc4f1d34a39178f9a4937ff35196f8bcdde2a67b5b5809c61fe"}, - {file = "cython-3.3.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bc2f2a6b65a991666cfd35a35bab0cd88ffba4df2f601edb6e76cc8116de24b9"}, - {file = "cython-3.3.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:23942b0662642927a55676e4b26e6840fb166dd7d76436384685227e7e8619a4"}, - {file = "cython-3.3.0-cp313-cp313-win_amd64.whl", hash = "sha256:ab24d1a4fb6aaf0b5b6fcd75a6d70255fbd3130fa78884c26991f8d5502616b5"}, - {file = "cython-3.3.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0deedc2e9a5a664e1adfa4c2d310aa7b54903e1a647c274b6c9213f77a02d637"}, - {file = "cython-3.3.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:46072c0d404616b5e652a63882c79cc3f8a1d62635a8692f56ed0e416a4dfed8"}, - {file = "cython-3.3.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:82f94565b6001bab8e31bf52a0911672910b5735910612a2c0f772c719670006"}, - {file = "cython-3.3.0-cp314-cp314-win_amd64.whl", hash = "sha256:51999fb834365721b6c7f689cf6e2ec7c8667aae783df9eb5e589c290a414d9c"}, - {file = "cython-3.3.0-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:596e8df019372a2cd417805015022d42cb8ee4e1803ccdc11ed00e451625fb66"}, - {file = "cython-3.3.0-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a36c34d1950845b8ac148653b07cdc62421a4b0d9abfcc849e69f1c4ff9919d"}, - {file = "cython-3.3.0-cp315-cp315-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b447f6906e0555f05dc4742ef1f99091b1e5d9aa9f16616e772fbf9ff6271616"}, - {file = "cython-3.3.0-cp315-cp315-win_amd64.whl", hash = "sha256:b55c72e8eccdd508c8de3cf3bbc543aafbb3bf6a518e1ee20358d3241cd780ef"}, - {file = "cython-3.3.0-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:e0d2713d2b292c826bc21dc8732bd9e47628103aa3764180c881e04b3fef95dc"}, - {file = "cython-3.3.0-cp39-abi3-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:169e56fd411f4cd5bba51c82f8239421d547a846099db2b261e4aed48ba9f51f"}, - {file = "cython-3.3.0-cp39-abi3-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:29f38ebafdf23e3da2516f40c4d065da38bfe002181bf93e2b8cf1262449aba6"}, - {file = "cython-3.3.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:75c4ae8a6d3a5ccf3cdaba8ab32e6a8d0cd38e3a476aa7ac12df8f8171a8d570"}, - {file = "cython-3.3.0-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:b94fb5613b9fe34c27d13ec9972dc0dcd2a2155db2902e93921cadc162610a38"}, - {file = "cython-3.3.0-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:c4558ba85849ab65dc57e10fd0efb13fabd9d3c09981a2566e18dec7cf47586a"}, - {file = "cython-3.3.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:311a016369adfd1e0015c4f9819168fc0e518451d7efb4435c30d65a3a26d52b"}, - {file = "cython-3.3.0-cp39-abi3-win32.whl", hash = "sha256:90869072e50b7c8904fe1dd7810321ae901fd5637a6eec6646ed9c57f9eb1081"}, - {file = "cython-3.3.0-cp39-abi3-win_arm64.whl", hash = "sha256:dce56c26d388f00a19426371b6926bf2f77c5c03b71d5273e4556c68be98c2dd"}, - {file = "cython-3.3.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:14e825253455e943ca765a95096b355745558436b0c46c24856de9269cc4dbd9"}, - {file = "cython-3.3.0-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:843d7134e784e7b320ef387512e89f1b29af80c641e176dfa8eabd52aab61c3c"}, - {file = "cython-3.3.0-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:26a5e536fc68e85a9de091a0b51c42c5ac834f8d00aaa43f227cbc3efa797ae5"}, - {file = "cython-3.3.0-cp39-cp39-win_amd64.whl", hash = "sha256:66d86b6a1548ae64851b211e3c3504535814b8c8e6c46ddcaf01062bf8d5fad2"}, - {file = "cython-3.3.0-py3-none-any.whl", hash = "sha256:9b24b5c8cd536946b62086fcafee6d5509d3f549f72d553d2336af87ffbe0da1"}, - {file = "cython-3.3.0.tar.gz", hash = "sha256:eed0d93fbca7087f143b42c34b05a825849bdf17f101572c2105acfa49aa88b8"}, -] - [[package]] name = "debugpy" version = "1.8.21" @@ -1015,37 +966,6 @@ files = [ {file = "distlib-0.4.3.tar.gz", hash = "sha256:f152097224a0ae24be5a0f6bae1b9359af82133bce63f98a95f86cae1aede9ed"}, ] -[[package]] -name = "distro" -version = "1.9.0" -description = "Distro - an OS platform information API" -optional = false -python-versions = ">=3.6" -groups = ["main"] -files = [ - {file = "distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2"}, - {file = "distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed"}, -] - -[[package]] -name = "eagerpy" -version = "0.30.0" -description = "EagerPy is a thin wrapper around PyTorch, TensorFlow Eager, JAX and NumPy that unifies their interface and thus allows writing code that works natively across all of them." -optional = false -python-versions = "*" -groups = ["main"] -files = [ - {file = "eagerpy-0.30.0-py3-none-any.whl", hash = "sha256:79c461b04577f02bf3b48191b2f911b55521204df99ec02288d96bfa34f13d80"}, - {file = "eagerpy-0.30.0.tar.gz", hash = "sha256:014c02b5a7f7e19f8471885cf8aa469f2e9cf518c88400f20b6b8db83d413106"}, -] - -[package.dependencies] -numpy = "*" -typing-extensions = ">=3.7.4.1" - -[package.extras] -testing = ["pytest (>=5.3.5)", "pytest-cov (>=2.8.1)"] - [[package]] name = "executing" version = "2.2.1" @@ -1189,30 +1109,6 @@ type1 = ["xattr ; sys_platform == \"darwin\""] unicode = ["unicodedata2 (>=17.0.0) ; python_version <= \"3.14\""] woff = ["brotli (>=1.0.1) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; platform_python_implementation != \"CPython\"", "zopfli (>=0.1.4)"] -[[package]] -name = "foolbox" -version = "3.3.4" -description = "Foolbox is an adversarial attacks library that works natively with PyTorch, TensorFlow and JAX" -optional = false -python-versions = "*" -groups = ["main"] -files = [ - {file = "foolbox-3.3.4-py3-none-any.whl", hash = "sha256:11049d515d38e765e206a73ace9cc648b81a0d6d1da9ad9057e149ee117989cb"}, - {file = "foolbox-3.3.4.tar.gz", hash = "sha256:1276cb1c1f636d1e6db08fb0d5cb00ec02144f145465166192dc86a86fce9e1c"}, -] - -[package.dependencies] -eagerpy = ">=0.30.0" -GitPython = ">=3.0.7" -numpy = "*" -requests = ">=2.24.0" -scipy = "*" -setuptools = "*" -typing-extensions = ">=3.7.4.1" - -[package.extras] -testing = ["pytest (>=7.1.1)", "pytest-cov (>=3.0.0)"] - [[package]] name = "fqdn" version = "1.5.1" @@ -1265,40 +1161,6 @@ test-downstream = ["aiobotocore (>=2.5.4,<3.0.0)", "dask[dataframe,test]", "moto test-full = ["adlfs", "aiohttp (!=4.0.0a0,!=4.0.0a1)", "backports-zstd ; python_version < \"3.14\"", "cloudpickle", "dask", "distributed", "dropbox", "dropboxdrivefs", "fastparquet", "fusepy", "gcsfs (>=2026.4.0)", "jinja2", "kerchunk", "libarchive-c", "lz4", "notebook", "numpy", "ocifs", "pandas (<3.0.0)", "panel", "paramiko", "pyarrow (>=1)", "pyftpdlib", "pygit2", "pytest", "pytest-asyncio (!=0.22.0)", "pytest-benchmark", "pytest-cov", "pytest-mock", "pytest-recording", "pytest-rerunfailures", "python-snappy", "requests", "s3fs (>=2026.6.0)", "smbprotocol", "tqdm", "urllib3", "zarr (<3.2.0)", "zstandard ; python_version < \"3.14\""] tqdm = ["tqdm"] -[[package]] -name = "gitdb" -version = "4.0.12" -description = "Git Object Database" -optional = false -python-versions = ">=3.7" -groups = ["main"] -files = [ - {file = "gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf"}, - {file = "gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571"}, -] - -[package.dependencies] -smmap = ">=3.0.1,<6" - -[[package]] -name = "gitpython" -version = "3.1.61" -description = "GitPython is a Python library used to interact with Git repositories" -optional = false -python-versions = ">=3.7" -groups = ["main"] -files = [ - {file = "gitpython-3.1.61-py3-none-any.whl", hash = "sha256:8ab28c9da863cdd9e7d7694ec46cf3e6c9a12d8a30a1acd3447aec11975d530c"}, - {file = "gitpython-3.1.61.tar.gz", hash = "sha256:f51c24d8c0f733a195447385f5774a5dfe8767f5acfd7994a33755644c6ecc95"}, -] - -[package.dependencies] -gitdb = ">=4.0.1,<5" - -[package.extras] -doc = ["sphinx (>=7.4.7,<8)", "sphinx-autodoc-typehints", "sphinx_rtd_theme"] -test = ["basedpyright (==1.39.9) ; python_version >= \"3.9\" and sys_platform != \"cygwin\"", "coverage[toml]", "ddt (>=1.1.1,!=1.4.3)", "mock ; python_version < \"3.8\"", "mypy (==1.18.2) ; python_version >= \"3.9\"", "pre-commit", "pytest (>=7.3.1)", "pytest-cov", "pytest-instafail", "pytest-mock", "pytest-sugar", "typing-extensions ; python_version < \"3.11\""] - [[package]] name = "gymnasium" version = "1.3.0" @@ -1411,7 +1273,7 @@ version = "3.19" description = "Internationalized Domain Names in Applications (IDNA)" optional = false python-versions = ">=3.9" -groups = ["main", "dev"] +groups = ["dev"] files = [ {file = "idna-3.19-py3-none-any.whl", hash = "sha256:815e7be7a7806d54abb586dc943addc79e8b2ee16915059658cbeff4b1b43bf4"}, {file = "idna-3.19.tar.gz", hash = "sha256:5e0811a4383b21dc5838069f801c4fb62113b7447663d2530d2bd6e77b49bf15"}, @@ -2167,42 +2029,6 @@ dev = ["changelist (==0.5)", "spin (==0.15)"] lint = ["pre-commit (==4.3.0)"] test = ["coverage[toml] (>=7.2)", "pytest (>=8.0)", "pytest-cov (>=5.0)"] -[[package]] -name = "llvmlite" -version = "0.49.0" -description = "lightweight wrapper around basic LLVM functionality" -optional = false -python-versions = ">=3.10" -groups = ["main"] -files = [ - {file = "llvmlite-0.49.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:ee81e96c15a6f870918f1eb60c913551c16aa23defb4f5f1acfa660d6a0aaac2"}, - {file = "llvmlite-0.49.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:854941c2267fd4fc5b2ce02b8af8ecdffa79fb7784591d3a89370322039ea09f"}, - {file = "llvmlite-0.49.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:da7b64474ac15ca595efa2644d5c6836638ccf70709fad3aba3fc56a55966928"}, - {file = "llvmlite-0.49.0-cp310-cp310-win_amd64.whl", hash = "sha256:b352c14353330c879e339b8f8d7491d565fe94242697714a24e80bd757202384"}, - {file = "llvmlite-0.49.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:4b0e710880b7cc910392bd6b9f1bbf468fed99b182e4420d51598f36114b3dce"}, - {file = "llvmlite-0.49.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a8c0fc9d624bdc30a3d2db11eb2fb98f80fb209d20b37604eda516cd9b699cf4"}, - {file = "llvmlite-0.49.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:20496a5c9fdb8179fb9300e7d19f6782555d98aeeb4a322264aa7fd99f980618"}, - {file = "llvmlite-0.49.0-cp311-cp311-win_amd64.whl", hash = "sha256:6a5b06c1b5fc4ae4c9b169b065f42b719448ef1f873687ef224ef69969b75ec3"}, - {file = "llvmlite-0.49.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:b541c8fac3450db7574d1f53cf9dff83f285bfed9d69bf81fe71fc2a7d4f97fe"}, - {file = "llvmlite-0.49.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6acba646d88abbc87d5c113a3d62c1fbf8b8fee11c6493f516803e30f21ae870"}, - {file = "llvmlite-0.49.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4ec8ad805e7515cb8440a690eb3cef4d34acb29eef80b705ec4e1c1ad3c43c68"}, - {file = "llvmlite-0.49.0-cp312-cp312-win_amd64.whl", hash = "sha256:3a9c9e3af4e214acfefa4f73ebe7bc3fb35854a62b654edb3953f5ae33c08ba3"}, - {file = "llvmlite-0.49.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:039fa4054a06f537fb39248d4472284ca96be311a142ec09e69f95630ab469cc"}, - {file = "llvmlite-0.49.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ddc7aecd4f56397ed6e8f120ec5dcd5a1a8f0e6032ca4af413462792d4dca2e3"}, - {file = "llvmlite-0.49.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d3dee64784201b64c13a8df62c48a4f4218858faaa65889866bb29bdc243c038"}, - {file = "llvmlite-0.49.0-cp313-cp313-win_amd64.whl", hash = "sha256:a1b414dc6b164738ec39dd8987cea73829057b7dd92fc6d91b52838385fc1dd2"}, - {file = "llvmlite-0.49.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:80a84683d04516bb51da1bbeebddaf2c2f558809c93078a8f91807909ae331f8"}, - {file = "llvmlite-0.49.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4281a0171d66d2098adce4ba706b8c550b1b10718650f682d64cde16e84e4de5"}, - {file = "llvmlite-0.49.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b095f15fb12c4d90495df5b1a3772b4732cc408398b204a787dbedd370e09c69"}, - {file = "llvmlite-0.49.0-cp314-cp314-win_amd64.whl", hash = "sha256:294e2f0b70aef8f92d0ae7b203e2609f08beb39437eee73de59a21669331aae9"}, - {file = "llvmlite-0.49.0-cp314-cp314-win_arm64.whl", hash = "sha256:95d1071023ed858b79f6971954fd7cc1f5dbcbab987718a4ccbe1411e47d0b81"}, - {file = "llvmlite-0.49.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:f3f2ff0aeb17d34fcce9f79b99baac441cfd3efa41b83e233ca4530a72381f72"}, - {file = "llvmlite-0.49.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d5555ea1d63928481cbf7fcb1d67452b216c7e5b393a4eb7aa1401e67f2a4fc4"}, - {file = "llvmlite-0.49.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e32adb84fdaae28aeb86fdb6253084ee707ee157289a2e98fe3caf48a62bee82"}, - {file = "llvmlite-0.49.0-cp314-cp314t-win_amd64.whl", hash = "sha256:be637e465010bc9c50f070468f7f1cf5385e92fee364d192dd5e6cea790ecba9"}, - {file = "llvmlite-0.49.0.tar.gz", hash = "sha256:00f16db782f4a13c78c5804aedc434e46794a77e89999a168f9401106270e50a"}, -] - [[package]] name = "markupsafe" version = "3.0.3" @@ -2627,46 +2453,6 @@ jupyter-server = ">=1.8,<3" [package.extras] test = ["pytest", "pytest-console-scripts", "pytest-jupyter", "pytest-tornasync"] -[[package]] -name = "numba" -version = "0.67.0" -description = "compiling Python code using LLVM" -optional = false -python-versions = ">=3.10" -groups = ["main"] -files = [ - {file = "numba-0.67.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:8c0e88acd4341ddf40779db3c0228b9188aca7fcab5f5f3ce9949a1fc71e9a02"}, - {file = "numba-0.67.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d6c8e9ba3f9602471e8c6f563ffcce8db8046741f0bafb782a052e41dc6b6861"}, - {file = "numba-0.67.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:694c81c6560b2b47e5fc1dc39c29175b907adf862d9af0af801453400a022a61"}, - {file = "numba-0.67.0-cp310-cp310-win_amd64.whl", hash = "sha256:ed333e0af4386294e7f03e550e01411856b6935e717d859225e0a7338c6b6795"}, - {file = "numba-0.67.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:3fa3d1b27f96f2c0d54513d953d7197886aa1eaa7d2439a0eedc44d993fb181a"}, - {file = "numba-0.67.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8c80c847301dc33dc8f84a97a952004023d9a05578ae4512b087176264cc1960"}, - {file = "numba-0.67.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e7a7b0121466f1e9a8a074b0545fe90e16389623abf979b5d7c299dca1294d7e"}, - {file = "numba-0.67.0-cp311-cp311-win_amd64.whl", hash = "sha256:cfba1ac34f0363fb1a250a10e97240780d11e05227892f7286b26fbfd0ad58ce"}, - {file = "numba-0.67.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:6004d8d5f28d4028687fb2d972d629295b13685943bd2ed5cd8810c3b848e219"}, - {file = "numba-0.67.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f63d43db06b4756424d6d2484737c902e0ae944a0eec3e8b0b4de2c695b15caa"}, - {file = "numba-0.67.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:76d3335aaeffb9dc88309420890e73497a00be08a7530441bc2b58ffe025bfa5"}, - {file = "numba-0.67.0-cp312-cp312-win_amd64.whl", hash = "sha256:50e2b72406c18cda5dd7431b0082cb85ea94e06c64c33607248fc8bef92cfb81"}, - {file = "numba-0.67.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:77e1c7173fee57a0d84e006c7e70346689d6cb3e7db503489bae58646b4eff7b"}, - {file = "numba-0.67.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9c4953387c77864b596d8296e2cfbdef82b0eea4166ab4864b05d226c51143e0"}, - {file = "numba-0.67.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:88f6e0f5cb6c545e158b6ef0496c01b6d6958a7ccc6634a1576a94bbbab29ff2"}, - {file = "numba-0.67.0-cp313-cp313-win_amd64.whl", hash = "sha256:b68ad5125fe245339cc8dcc036081fc1ea482c5063387b9612a76ccd83dc91cd"}, - {file = "numba-0.67.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:f99f880ff25f418a67f9a1d00d0ddfbc63430f627b523e515085a592a7567f4b"}, - {file = "numba-0.67.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5269245a675abdd3e2c35ec6bb2f250355effa9032514d8f2354f0d2d10854bd"}, - {file = "numba-0.67.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f074a8e23db78490f11a3930c940be758316c10ac5985be83d2f298dc080acf7"}, - {file = "numba-0.67.0-cp314-cp314-win_amd64.whl", hash = "sha256:4d576e62bf2c9370f61312b51573c4bb1f3fe96798bbab56730847a368a316c4"}, - {file = "numba-0.67.0-cp314-cp314-win_arm64.whl", hash = "sha256:7930748ce8355d2a5a28602abab056a61fdc676d17377f27d17993905428171f"}, - {file = "numba-0.67.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:4a2ed006635bbd0fe45681ed49f3b4f4bad1abf0c233bcc5842c9e3a34cabd61"}, - {file = "numba-0.67.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:aa5f002f665bec321b950dacaa26ee009e1d720f6ac9d9856eed5efe1caa03a6"}, - {file = "numba-0.67.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:83ab968b0e0fa744eba03351282dd8000796e6ec8e4518f47bd3ed86c0a20c7b"}, - {file = "numba-0.67.0-cp314-cp314t-win_amd64.whl", hash = "sha256:00c964a5b94d3ae82d83ac162cd610755875b98dadb779fdde06e6bfcdbca47e"}, - {file = "numba-0.67.0.tar.gz", hash = "sha256:cd75aa535b33fa05d9d930b1ae8af9f97a2881e96d72dfb38ec9b78284d9f851"}, -] - -[package.dependencies] -llvmlite = "==0.49.*" -numpy = ">=1.22,<2.6" - [[package]] name = "numpy" version = "2.4.6" @@ -3881,7 +3667,7 @@ version = "2.34.2" description = "Python HTTP for Humans." optional = false python-versions = ">=3.10" -groups = ["main", "dev"] +groups = ["dev"] files = [ {file = "requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0"}, {file = "requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed"}, @@ -4068,30 +3854,6 @@ files = [ {file = "rpds_py-2026.6.3.tar.gz", hash = "sha256:1cebd1337c242e4ec2293e541f712b2da849b29f48f0c293684b71c0632625d4"}, ] -[[package]] -name = "scikit-build" -version = "0.19.1" -description = "Improved build system generator for Python C/C++/Fortran/Cython extensions" -optional = false -python-versions = ">=3.8" -groups = ["main"] -files = [ - {file = "scikit_build-0.19.1-py3-none-any.whl", hash = "sha256:a4fa78a7222ca40d5e8bae6700514bbf7efea4346efd9fc1f8cfca837cd1de71"}, - {file = "scikit_build-0.19.1.tar.gz", hash = "sha256:b9a8d07fca2d5d10d93220bc57a685161d72af1fc76285d55c564ddaa862e584"}, -] - -[package.dependencies] -distro = "*" -packaging = "*" -setuptools = ">=42.0.0" -wheel = ">=0.32.0" - -[package.extras] -cov = ["coverage[toml] (>=4.2)", "pytest-cov (>=2.7.1)"] -docs = ["pygments", "sphinx (>=4)", "sphinx-issues", "sphinx-rtd-theme (>=1.0)", "sphinxcontrib-moderncmakedomain (>=3.19)"] -doctest = ["ubelt (>=0.8.2)", "xdoctest (>=0.10.0)"] -test = ["build (>=0.7)", "cython (>=0.25.1)", "pip", "pytest (>=6.0.0)", "pytest-mock (>=1.10.4)", "requests", "virtualenv"] - [[package]] name = "scikit-image" version = "0.26.0" @@ -4359,18 +4121,6 @@ files = [ {file = "six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81"}, ] -[[package]] -name = "smmap" -version = "5.0.3" -description = "A pure Python implementation of a sliding window memory map manager" -optional = false -python-versions = ">=3.7" -groups = ["main"] -files = [ - {file = "smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f"}, - {file = "smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c"}, -] - [[package]] name = "soupsieve" version = "2.9.2" @@ -4850,7 +4600,7 @@ version = "2.7.0" description = "HTTP library with thread-safe connection pooling, file post, and more." optional = false python-versions = ">=3.10" -groups = ["main", "dev"] +groups = ["dev"] files = [ {file = "urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897"}, {file = "urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c"}, @@ -4937,22 +4687,7 @@ docs = ["Sphinx (>=6.0)", "myst-parser (>=2.0.0)", "sphinx_rtd_theme (>=1.1.0)"] optional = ["python-socks", "wsaccel"] test = ["pytest", "websockets"] -[[package]] -name = "wheel" -version = "0.48.0" -description = "Command line tool for manipulating wheel files" -optional = false -python-versions = ">=3.9" -groups = ["main"] -files = [ - {file = "wheel-0.48.0-py3-none-any.whl", hash = "sha256:3217dcc807155e45db462d7ef2431f5ddda0d7273b700d05a67b271ceb1287ab"}, - {file = "wheel-0.48.0.tar.gz", hash = "sha256:94800765601e9171bf5d58d066e640662842bcedcbab982b2c90787a2c987322"}, -] - -[package.dependencies] -packaging = ">=24.0" - [metadata] lock-version = "2.1" python-versions = ">=3.11,<3.14" -content-hash = "d2b2b82985da02a1edc6c58c7d7e14baa052c80fcc6bc7eec09a69a22abb0067" +content-hash = "634a5f974068c11e11ab0a6a7ed03785e62ce36116d6b911ea590e39e8d86c5c" diff --git a/pyproject.toml b/pyproject.toml index 35e05070a..9d5844aad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,44 +1,61 @@ -[tool.poetry] +[project] name = "bindsnet" version = "0.3.4" description = "Spiking neural networks for ML in Python" -authors = [ "Hananel Hazan ", "Daniel Saunders", "Darpan Sanghavi", "Hassaan Khan" ] +authors = [ + { name = "Hananel Hazan", email = "hananel@hazan.org.il" }, + { name = "Daniel Saunders" }, + { name = "Darpan Sanghavi" }, + { name = "Hassaan Khan" }, +] license = "AGPL-3.0-only" +license-files = ["LICENSE"] readme = "README.md" -repository = "https://github.com/BindsNET/bindsnet" -documentation = "https://bindsnet-docs.readthedocs.io/" +requires-python = ">=3.11,<3.14" keywords = ["spiking", "neural", "networks", "pytorch"] +classifiers = [ + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", +] +dependencies = [ + "numpy>=2,<3", + "scipy>=1,<2", + "torch>=2.14,<3", + "torchvision>=0.29,<1", + "tensorboardX>=2.6.4,<3", + "tqdm>=4,<5", + "matplotlib>=3,<4", + "ale-py>=0,<1", + "gymnasium[atari]>=1,<2", + "scikit-image>=0,<1", + "scikit-learn>=1,<2", + "opencv-python>=4.14,<6", + "pandas>=3,<4", +] + +[project.urls] +Repository = "https://github.com/BindsNET/bindsnet" +Documentation = "https://bindsnet-docs.readthedocs.io/" +Changelog = "https://github.com/BindsNET/bindsnet/blob/master/CHANGELOG.md" +# Poetry-only detail: on Linux and Windows, install torch/torchvision from the CUDA 13.0 +# wheel index. PyPI users get the torch build that PyPI serves. [tool.poetry.dependencies] -python = ">=3.11,<3.14" -numpy = "^2" -numba = "^0" -scipy = "^1" -Cython = "^3" torch = [ - {version = "2.14.0", markers = "sys_platform != 'darwin'", source = "torch+cu130"}, - {version = "2.14.0", markers = "sys_platform == 'darwin'" }, + { markers = "sys_platform != 'darwin'", source = "torch+cu130" }, + { markers = "sys_platform == 'darwin'" }, ] torchvision = [ - {version = "0.29.0", markers = "sys_platform != 'darwin'", source = "torch+cu130"}, - {version = "0.29.0", markers = "sys_platform == 'darwin'" }, + { markers = "sys_platform != 'darwin'", source = "torch+cu130" }, + { markers = "sys_platform == 'darwin'" }, ] -tensorboardX = "^2.6.4" -tqdm = "^4" -matplotlib = "^3" -ale-py = "^0" -gymnasium = {extras = ["atari"], version = "^1"} -scikit-build = "^0" -scikit-image = "^0" -scikit-learn = "^1" -opencv-python = ">=4.14,<6" -pandas = "^3" -foolbox = "^3" [[tool.poetry.source]] - name = "torch+cu130" - url = "https://download.pytorch.org/whl/cu130" - priority = "explicit" +name = "torch+cu130" +url = "https://download.pytorch.org/whl/cu130" +priority = "explicit" [tool.poetry.group.dev.dependencies] pytest = "^9" @@ -50,7 +67,7 @@ black = "^26" autoflake = "^2" [build-system] -requires = ["setuptools", "poetry-core>=1.0.0"] +requires = ["poetry-core>=2.0"] build-backend = "poetry.core.masonry.api" [tool.isort] diff --git a/setup.py b/setup.py deleted file mode 100644 index bac24a43d..000000000 --- a/setup.py +++ /dev/null @@ -1,6 +0,0 @@ -#!/usr/bin/env python - -import setuptools - -if __name__ == "__main__": - setuptools.setup() diff --git a/test/conversion/test_conversion.py b/test/conversion/test_conversion.py index 8f5deb079..5eafdb8d2 100644 --- a/test/conversion/test_conversion.py +++ b/test/conversion/test_conversion.py @@ -2,7 +2,7 @@ import torch.nn as nn import torch.nn.functional as F -from bindsnet.conversion import ann_to_snn +from bindsnet.conversion import ann_to_snn, data_based_normalization class FullyConnectedNetwork(nn.Module): @@ -36,6 +36,38 @@ def test_conversion_2(): snn = ann_to_snn(ann, data=data, input_shape=(784,)) +def test_conversion_from_path(tmp_path): + """ + ``ann_to_snn`` accepts a path to a saved network. + + Regression test: this path called ``torch.load`` without + ``weights_only=False``, so it broke outright when PyTorch 2.6 changed that + default to ``True``. Only the in-memory ``nn.Module`` form was tested, so + the breakage went unnoticed. + """ + ann = FullyConnectedNetwork() + file_path = str(tmp_path / "ann.pt") + torch.save(ann, file_path) + + snn = ann_to_snn(file_path, input_shape=(784,)) + + assert snn is not None + + +def test_data_based_normalization_from_path(tmp_path): + """ + ``data_based_normalization`` accepts a path to a saved network. Same + PyTorch 2.6 regression as ``test_conversion_from_path``. + """ + ann = FullyConnectedNetwork() + file_path = str(tmp_path / "ann.pt") + torch.save(ann, file_path) + + normalized = data_based_normalization(file_path, data=torch.rand(20, 784)) + + assert isinstance(normalized, nn.Module) + + def main(): test_conversion_1() test_conversion_2() diff --git a/test/datasets/test_cache_serialization.py b/test/datasets/test_cache_serialization.py new file mode 100644 index 000000000..67b0185ad --- /dev/null +++ b/test/datasets/test_cache_serialization.py @@ -0,0 +1,52 @@ +import numpy as np +import torch + + +class TestSpokenMNISTCache: + """ + Pins the serialization contract of the ``SpokenMNIST`` processed-data cache. + + ``SpokenMNIST`` writes its processed data with + ``torch.save((audio, labels), ...)`` and reads it back with + ``torch.load(..., weights_only=True)``. ``weights_only=True`` refuses to + execute code while loading, which is only usable because this cache holds + nothing but tensors. These tests pin that property, so that if the cached + payload ever gains a non-tensor object the failure shows up here rather + than as a broken dataset load for a user. + + The real loader needs a download, so these tests exercise the same + save/load pair on the same payload shape instead: ``audio`` is a list of + 2-D float tensors (filter banks, one per utterance) and ``labels`` is a + 1-D float tensor, per ``SpokenMNIST.process_data``. + """ + + def test_cache_round_trips_under_weights_only(self, tmp_path): + audio = [torch.rand(7, 13), torch.rand(11, 13)] + labels = torch.Tensor([3.0, 8.0]) + + path = str(tmp_path / "audio.pt") + torch.save((audio, labels), open(path, "wb")) + + _audio, _labels = torch.load(open(path, "rb"), weights_only=True) + + assert len(_audio) == len(audio) + for loaded, original in zip(_audio, audio): + assert torch.equal(loaded, original) + assert torch.equal(_labels, labels) + + def test_numpy_derived_filter_banks_round_trip(self, tmp_path): + """ + ``process_data`` builds filter banks through NumPy before they become + tensors. This pins that the converted result still loads under + ``weights_only=True``. + """ + audio = [torch.Tensor(np.random.rand(5, 13).astype(np.float32))] + labels = torch.Tensor([1.0]) + + path = str(tmp_path / "audio.pt") + torch.save((audio, labels), open(path, "wb")) + + _audio, _labels = torch.load(open(path, "rb"), weights_only=True) + + assert torch.equal(_audio[0], audio[0]) + assert torch.equal(_labels, labels) diff --git a/test/network/test_learning.py b/test/network/test_learning.py index 69eee8493..33c47608d 100644 --- a/test/network/test_learning.py +++ b/test/network/test_learning.py @@ -376,7 +376,9 @@ def test_mstdp_reset_clears_fast_path_spike_lag(self): assert rule._prev_source_s is None assert rule._prev_target_s is None - @pytest.mark.parametrize("rule", [mcc.MSTDP, mcc.MSTDPET, mcc.PostPre]) + @pytest.mark.parametrize( + "rule", [mcc.MSTDP, mcc.MSTDPET, mcc.PostPre, mcc.Hebbian, mcc.DiehlAndCook] + ) def test_reset_before_first_run_does_not_raise(self, rule): # Some of this state is built lazily on the first update, because only # then are the batch size and device known. Resetting a network before @@ -399,7 +401,9 @@ def test_postpre_reset_clears_average_buffers(self): assert rule.average_buffer_index_pre == 0 assert rule.average_buffer_index_post == 0 - @pytest.mark.parametrize("rule", [mcc.MSTDP, mcc.MSTDPET, mcc.PostPre]) + @pytest.mark.parametrize( + "rule", [mcc.MSTDP, mcc.MSTDPET, mcc.PostPre, mcc.Hebbian, mcc.DiehlAndCook] + ) def test_episodes_are_independent_after_reset(self, rule): # The symptom #777 reported: with a reset between them, two identical # episodes must produce identical weights. Before the fix the second diff --git a/test/network/test_learning_rule_specs.py b/test/network/test_learning_rule_specs.py new file mode 100644 index 000000000..7e64b547b --- /dev/null +++ b/test/network/test_learning_rule_specs.py @@ -0,0 +1,686 @@ +# language=rst +""" +Validation of the weight-modifying learning rules against the equations of the +papers they implement. Each reference below is a from-scratch transcription of +the cited equations, driven with the spike trains actually produced by the +network, so the tests check the rule (not the neuron model) step by step. + +Sources +------- +* Pair-based trace STDP (``PostPre``, ``WeightDependentPostPre``, ``Hebbian``): + Morrison, Diesmann & Gerstner (2008), *Biol. Cybern.* 98:459-478, + Sect. 4.1, eqs. (11)-(15): + + .. math:: + + \\dot x_j = -x_j/\\tau_x + \\sum_f \\delta(t - t_j^f) \\qquad (11) + + \\dot y_i = -y_i/\\tau_y + \\sum_f \\delta(t - t_i^f) \\qquad (12) + + \\Delta w_{ij}^+(t_i^f) = F_+(w_{ij})\\, x_j(t_i^f) \\qquad (13) + + \\Delta w_{ij}^-(t_j^f) = -F_-(w_{ij})\\, y_i(t_j^f) \\qquad (14) + + with ``j`` presynaptic (``source``) and ``i`` postsynaptic (``target``). + ``PostPre`` is the additive case :math:`F_+ = \\nu_\\text{post}`, + :math:`F_- = \\nu_\\text{pre}`; ``WeightDependentPostPre`` is the soft-bounded + (multiplicative) case :math:`F_+ = \\nu_\\text{post}(w_\\max - w)`, + :math:`F_- = \\nu_\\text{pre}(w - w_\\min)`. Traces follow Sect. 2.3: the + accumulating trace (``traces_additive=True``) adds 1 per spike; the saturating + trace with :math:`A = 1` (``traces_additive=False``) resets to 1 on each spike. + ``Hebbian`` is the same trace machinery with both terms positive (BindsNET's + own definition; no paper equation). + +* Diehl & Cook STDP (``DiehlAndCook``): Diehl & Cook (2015), *Front. Comput. + Neurosci.* 9:99, Sect. 2.3 "Learning": on each post-synaptic spike + + .. math:: + + \\Delta w = \\eta\\,(x_\\text{pre} - x_\\text{tar})\\,(w_\\max - w)^\\mu + + with the pre-synaptic trace increased by 1 per spike and decaying + exponentially; pre-synaptic spikes do not change the weight. The paper gives + no numeric values for :math:`x_\\text{tar}` or :math:`\\mu`. + +* Reward-modulated STDP (``MSTDP``, ``MSTDPET``): Florian (2007), *Neural + Comput.* 19:1468-1502, discrete-time eqs. (3.9)-(3.12) and (2.7)-(2.8). Those + rules are validated in ``test_mstdp_florian.py``; this file only records the + equation numbers and the timing convention: + :math:`w(t+\\delta t) = w(t) + \\gamma\\, r(t+\\delta t)\\, \\zeta(t)` (3.9), i.e. + the reward supplied at a step multiplies the eligibility of the *previous* + step (the default ``zero_lag=False``). + +* R-max (``Rmax``): Vasilaki, Fremaux, Urbanczik, Senn & Gerstner (2009), + *PLoS Comput. Biol.* 5(12):e1000586, eqs. (7), (8) and (13): + + .. math:: + + \\dot w_{ij} = \\alpha (R - b)\\, \\delta(t - t_\\text{hit})\\, e_{ij}(t) \\qquad (7) + + \\dot e_{ij} = -e_{ij}/\\tau_e + \\frac{g'}{g}\\Big[Y_i(t) - + \\frac{\\rho_i(t)}{1 + \\tau_c\\,\\rho_i(t)}\\Big] \\sum_f \\epsilon(t - t_j^f) + \\qquad (8) + + \\rho_i = g(u_i) = \\rho_0 \\exp\\big((u_i - u_0)/\\Delta u\\big) \\qquad (13) + + :math:`\\tau_c = 0` is the strict policy-gradient rule and + :math:`\\tau_c \\to \\infty` the naive Hebbian rule (Vasilaki et al., p. 3-4). + For the exponential :math:`g` of (13), :math:`g'/g = 1/\\Delta u` is a constant + absorbed into the learning rate (ibid.). BindsNET discretises (8) with a + forward-Euler decay and the per-step spike probability + :math:`p = 1 - e^{-\\rho\\,\\delta t}` in place of :math:`\\rho\\,\\delta t`. + +Spikes forced with the ``clamp`` argument of ``Network.run`` (and spikes removed with +``unclamp``) are applied inside ``Nodes.forward`` before the spike trace is updated, +so they are seen by the learning rules exactly like naturally generated spikes. +``TestClampEntersTraces`` pins that; the STDP-window test fires the post-synaptic +neuron through a teacher input and repeats it with ``clamp``. +""" + +import math + +import pytest +import torch + +from bindsnet.learning import ( + DiehlAndCook, + Hebbian, + PostPre, + Rmax, + WeightDependentPostPre, +) +from bindsnet.learning import MCC_learning +from bindsnet.models import DiehlAndCook2015 +from bindsnet.network import Network +from bindsnet.network.nodes import Input, LIFNodes, SRM0Nodes +from bindsnet.network.topology import Connection, MulticompartmentConnection +from bindsnet.network.topology_features import Weight + +TOL = 1e-5 + + +def _trace_step(trace, spikes, decay, additive): + """Morrison et al. (2008) Sect. 2.3 trace, one time step.""" + trace = trace * decay + if additive: + return trace + spikes + return torch.where(spikes > 0, torch.ones_like(trace), trace) + + +def _pair_stdp_reference( + pre, post, w0, nu_pre, nu_post, dt, tc_trace, additive, f, wmin, wmax +): + """ + Morrison et al. (2008) eqs. (11)-(14) in discrete time. ``pre``/``post`` are + ``[T, n]`` 0/1 spike arrays; ``f(w)`` returns ``(F_minus(w), F_plus(w))``. + Traces are updated before the weight change of the same step, matching the + network's order (``Nodes.forward`` then ``Connection.update``). + """ + w = w0.clone() + x = torch.zeros(pre.shape[1]) + y = torch.zeros(post.shape[1]) + decay = math.exp(-dt / tc_trace) + hist = [] + for t in range(pre.shape[0]): + x = _trace_step(x, pre[t], decay, additive) + y = _trace_step(y, post[t], decay, additive) + f_minus, f_plus = f(w) + w = w - nu_pre * f_minus * torch.outer(pre[t], y) # (14): pre spike + w = w + nu_post * f_plus * torch.outer(x, post[t]) # (13): post spike + w = w.clamp(wmin, wmax) # hard bounds, applied every step as in BindsNET + hist.append(w.clone()) + return torch.stack(hist) + + +def _run_pair_rule(rule, dt, additive, wmin, wmax, seed=0, T=40, n_in=12, n_out=6): + torch.manual_seed(seed) + net = Network(dt=dt) + net.add_layer( + Input(n=n_in, traces=True, traces_additive=additive, tc_trace=20.0), "in" + ) + net.add_layer( + LIFNodes(n=n_out, traces=True, traces_additive=additive, tc_trace=20.0), + "out", + ) + w0 = 0.5 * torch.rand(n_in, n_out) + conn = Connection( + net.layers["in"], + net.layers["out"], + w=w0.clone(), + nu=(1e-2, 3e-2), + update_rule=rule, + wmin=wmin, + wmax=wmax, + ) + net.add_connection(conn, "in", "out") + # Strong drive so the LIF layer actually spikes. + torch.manual_seed(seed + 1) + pre = torch.bernoulli(0.5 * torch.ones(T, n_in)) + net.layers["out"].thresh.fill_(-60.0) + w_hist, post_hist = [], [] + for t in range(T): + net.run(inputs={"in": pre[t : t + 1].byte()}, time=dt) + post_hist.append(net.layers["out"].s.view(-1).float().clone()) + w_hist.append(conn.w.detach().clone()) + post = torch.stack(post_hist) + assert post.sum() > 0, "post-synaptic layer never spiked; test is vacuous" + return w0, pre, post, torch.stack(w_hist), conn + + +class TestPairSTDPMorrison2008: + @pytest.mark.parametrize("dt", [1.0, 0.5]) + @pytest.mark.parametrize("additive", [False, True]) + def test_postpre_is_additive_pair_stdp(self, dt, additive): + w0, pre, post, w_b, conn = _run_pair_rule(PostPre, dt, additive, 0.0, 1.0) + w_ref = _pair_stdp_reference( + pre, + post, + w0, + 1e-2, + 3e-2, + dt, + 20.0, + additive, + lambda w: (1.0, 1.0), + 0.0, + 1.0, + ) + assert (w_b - w_ref).abs().max().item() < TOL + + @pytest.mark.parametrize("additive", [False, True]) + def test_weight_dependent_is_soft_bounded_pair_stdp(self, additive): + wmin, wmax = 0.0, 1.0 + w0, pre, post, w_b, conn = _run_pair_rule( + WeightDependentPostPre, 1.0, additive, wmin, wmax + ) + w_ref = _pair_stdp_reference( + pre, + post, + w0, + 1e-2, + 3e-2, + 1.0, + 20.0, + additive, + lambda w: (w - wmin, wmax - w), + wmin, + wmax, + ) + assert (w_b - w_ref).abs().max().item() < TOL + + def test_hebbian_both_terms_potentiate(self): + w0, pre, post, w_b, conn = _run_pair_rule(Hebbian, 1.0, False, 0.0, 1.0) + w = w0.clone() + x = torch.zeros(pre.shape[1]) + y = torch.zeros(post.shape[1]) + decay = math.exp(-1.0 / 20.0) + hist = [] + for t in range(pre.shape[0]): + x = _trace_step(x, pre[t], decay, False) + y = _trace_step(y, post[t], decay, False) + w = w + 1e-2 * torch.outer(pre[t], y) + 3e-2 * torch.outer(x, post[t]) + hist.append(w.clamp(0.0, 1.0).clone()) + w = w.clamp(0.0, 1.0) + assert (w_b - torch.stack(hist)).abs().max().item() < TOL + + @pytest.mark.parametrize("drive", ["teacher", "clamp"]) + def test_stdp_window_sign_and_shape(self, drive): + # Morrison (2008) eq. (10): pre-before-post potentiates by + # F_+ exp(-|dt|/tau_+); post-before-pre depresses by F_- exp(-|dt|/tau_-). + # The post neuron is fired either by a strong "teacher" input or by the + # ``clamp`` argument of ``Network.run``; both must give the same window. + tc = 20.0 + for delta in (1, 5, 15): + for pre_first in (True, False): + net = Network(dt=1.0) + net.add_layer(Input(n=1, traces=True, tc_trace=tc), "in") + net.add_layer(Input(n=1), "teacher") + net.add_layer(LIFNodes(n=1, traces=True, tc_trace=tc), "out") + conn = Connection( + net.layers["in"], + net.layers["out"], + w=torch.zeros(1, 1), # no drive from the plastic synapse + nu=(1.0, 1.0), + update_rule=PostPre, + wmin=-10.0, + wmax=10.0, + ) + net.add_connection(conn, "in", "out") + net.add_connection( + Connection( + net.layers["teacher"], + net.layers["out"], + w=torch.full((1, 1), 100.0), + ), + "teacher", + "out", + ) + T = 40 + pre = torch.zeros(T, 1, 1) + teach = torch.zeros(T, 1, 1) + if pre_first: + pre[10, 0, 0] = 1 + teach[10 + delta - 1, 0, 0] = 1 # post fires one step later + else: + teach[10 - 1, 0, 0] = 1 + pre[10 + delta, 0, 0] = 1 + post_times = [] + for t in range(T): + if drive == "teacher": + net.run( + inputs={"in": pre[t].byte(), "teacher": teach[t].byte()}, + time=1, + ) + else: + # Force the post spike at the step the teacher would + # have fired it (one step after the teacher spike). + force = ( + teach[t - 1, 0].bool() if t > 0 else torch.zeros(1).bool() + ) + net.run( + inputs={ + "in": pre[t].byte(), + "teacher": torch.zeros(1, 1).byte(), + }, + time=1, + clamp={"out": force}, + ) + if net.layers["out"].s.any(): + post_times.append(t) + assert post_times == [10 + delta if pre_first else 10], post_times + change = conn.w.item() + expected = math.exp(-delta / tc) * (1 if pre_first else -1) + assert abs(change - expected) < 1e-5, (delta, pre_first, change) + + def test_diehl_and_cook_2015_rule_is_not_postpre(self): + # Diehl & Cook (2015), Methods "Learning": weights change only on + # postsynaptic spikes, Delta w = eta (x_pre - x_tar) (w_max - w)^mu. + # BindsNET's ``PostPre`` (used by ``DiehlAndCook2015``) instead has no + # x_tar term and a depression term on presynaptic spikes. Pin that fact + # so the deviation stays documented: a lone presynaptic spike changes + # the weight under PostPre (it would not under Diehl & Cook's rule). + net = Network(dt=1.0) + net.add_layer(Input(n=1, traces=True), "in") + net.add_layer(LIFNodes(n=1, traces=True), "out") + conn = Connection( + net.layers["in"], + net.layers["out"], + w=torch.full((1, 1), 0.5), + nu=(1.0, 1.0), + update_rule=PostPre, + wmin=0.0, + wmax=1.0, + ) + net.add_connection(conn, "in", "out") + net.layers["out"].x.fill_(0.3) # a lingering post-synaptic trace + pre = torch.zeros(3, 1, 1) + pre[1, 0, 0] = 1 + net.run(inputs={"in": pre.byte()}, time=3) + assert conn.w.item() < 0.5 # depressed by the pre spike alone + + +class TestMulticompartmentRulesMatchClassic: + """The ``MCC_learning`` PostPre / Hebbian / DiehlAndCook must apply the same + equations as the classic rules at any ``dt`` (the MCC PostPre used to scale + its update by ``dt``; per-spike increments must not depend on the step).""" + + @pytest.mark.parametrize("dt", [1.0, 0.5]) + @pytest.mark.parametrize( + "rule_pair", + [ + (PostPre, MCC_learning.PostPre), + (Hebbian, MCC_learning.Hebbian), + (DiehlAndCook, MCC_learning.DiehlAndCook), + ], + ) + def test_same_weights_step_by_step(self, rule_pair, dt): + classic, mcc = rule_pair + rule_kw = {"x_tar": 0.3, "mu": 1.0} if classic is DiehlAndCook else {} + torch.manual_seed(0) + w0 = 0.5 * torch.rand(12, 6) + torch.manual_seed(1) + pre = torch.bernoulli(0.5 * torch.ones(40, 12)).byte() + + def build(use_mcc): + net = Network(dt=dt) + net.add_layer(Input(n=12, traces=True, traces_additive=True), "in") + net.add_layer(LIFNodes(n=6, traces=True, thresh=-60.0), "out") + if use_mcc: + conn = MulticompartmentConnection( + net.layers["in"], + net.layers["out"], + device="cpu", + pipeline=[ + Weight( + "w", + w0.clone(), + range=[0.0, 1.0], + nu=(1e-2, 3e-2), + learning_rule=mcc, + **rule_kw, + ) + ], + ) + else: + conn = Connection( + net.layers["in"], + net.layers["out"], + w=w0.clone(), + nu=(1e-2, 3e-2), + update_rule=classic, + wmin=0.0, + wmax=1.0, + **rule_kw, + ) + net.add_connection(conn, "in", "out") + return net, conn + + net_a, conn_a = build(False) + net_b, conn_b = build(True) + for t in range(40): + net_a.run(inputs={"in": pre[t : t + 1]}, time=dt) + net_b.run(inputs={"in": pre[t : t + 1]}, time=dt) + wa, wb = conn_a.w, conn_b.pipeline[0].value + assert (wa - wb).abs().max().item() < TOL, t + assert (conn_a.w - w0).abs().max().item() > 1e-3, "vacuous: no learning" + + def test_mcc_postpre_update_independent_of_dt(self): + # One post spike with a pre trace of 1 must add exactly nu_post, + # whatever the simulation step. + for dt in (1.0, 0.5, 0.1): + net = Network(dt=dt) + net.add_layer(Input(n=1, traces=True), "in") + net.add_layer(LIFNodes(n=1, traces=True), "out") + conn = MulticompartmentConnection( + net.layers["in"], + net.layers["out"], + device="cpu", + pipeline=[ + Weight( + "w", + torch.full((1, 1), 0.5), + range=[0.0, 1.0], + nu=(0.0, 0.1), + learning_rule=MCC_learning.PostPre, + ) + ], + ) + net.add_connection(conn, "in", "out") + net.layers["in"].x.fill_(1.0) + net.layers["out"].s.fill_(True) + conn.update(learning=True) + assert conn.pipeline[0].value.item() == pytest.approx(0.6), dt + + +class TestRmaxVasilaki2009: + @staticmethod + def _build(tc_c, dt=1.0, seed=0, n_in=10, n_out=4): + torch.manual_seed(seed) + net = Network(dt=dt) + net.add_layer( + Input(n=n_in, traces=True, traces_additive=True, tc_trace=10.0), "in" + ) + net.add_layer(SRM0Nodes(n=n_out, tc_decay=10.0, thresh=-55.0), "out") + conn = Connection( + net.layers["in"], + net.layers["out"], + w=2.0 * torch.rand(n_in, n_out), + nu=1e-2, + update_rule=Rmax, + wmin=-10.0, + wmax=10.0, + tc_c=tc_c, + tc_e_trace=25.0, + ) + net.add_connection(conn, "in", "out") + return net, conn + + def test_srm0_escape_rate_eq13(self): + # rho = rho_0 exp((u - theta)/Delta u), spike probability 1 - exp(-rho dt). + net, conn = self._build(5.0, dt=0.5) + layer = net.layers["out"] + torch.manual_seed(3) + net.run( + inputs={"in": torch.bernoulli(0.5 * torch.ones(5, 10)).byte()}, + time=2.5, + reward=0.0, + ) + rho = layer.rho_0 * torch.exp((layer.v - layer.thresh) / layer.d_thresh) + assert torch.allclose(layer.rho, rho) + assert torch.allclose(layer.s_prob, 1.0 - torch.exp(-rho * 0.5)) + + @pytest.mark.parametrize("tc_c", [0.0, 5.0, 1e9]) + @pytest.mark.parametrize("dt", [1.0, 0.5]) + def test_eligibility_and_update_eq7_eq8(self, tc_c, dt): + net, conn = self._build(tc_c, dt=dt) + src, tgt = net.layers["in"], net.layers["out"] + w_ref = conn.w.detach().clone() + e = torch.zeros_like(w_ref) + torch.manual_seed(4) + T = 30 + pre = torch.bernoulli(0.5 * torch.ones(T, 10)).byte() + rewards = torch.randn(T) + for t in range(T): + net.run(inputs={"in": pre[t : t + 1]}, time=dt, reward=rewards[t].item()) + # Observed post-synaptic factors of this step (Y_i and rho_i dt as p_i). + Y = tgt.s.view(-1).float() + p = tgt.s_prob.view(-1) + eps = src.x.view(-1) # sum_f epsilon(t - t_j^f): additive pre trace + # Eq. (8), forward Euler in dt, with rho dt -> p and tau_c rho -> (tau_c/dt) p. + post_factor = Y - p / (1.0 + (tc_c / dt) * p) + e = e * (1.0 - dt / 25.0) + torch.outer(eps, post_factor) + # Eq. (7): the reward of this step gates the eligibility of this step. + w_ref = (w_ref + 1e-2 * rewards[t] * e).clamp(-10.0, 10.0) + assert (conn.w - w_ref).abs().max().item() < TOL, t + + def test_tc_c_limits(self): + # tau_c = 0: post factor is Y - p (policy gradient, mean-zero in + # expectation); tau_c -> inf: post factor is Y (naive Hebbian). + net, conn = self._build(0.0) + rule = conn.update_rule + tgt = net.layers["out"] + tgt.s = torch.tensor([[1, 0, 1, 0]], dtype=torch.bool) + tgt.s_prob = torch.tensor([[0.2, 0.4, 0.6, 0.8]]) + net.layers["in"].x = torch.ones(1, 10) + conn.update(reward=1.0, learning=True) + expected = torch.tensor([1.0, 0.0, 1.0, 0.0]) - tgt.s_prob.view(-1) + assert torch.allclose(rule.eligibility_trace[0], expected) + + net, conn = self._build(1e12) + rule = conn.update_rule + tgt = net.layers["out"] + tgt.s = torch.tensor([[1, 0, 1, 0]], dtype=torch.bool) + tgt.s_prob = torch.tensor([[0.2, 0.4, 0.6, 0.8]]) + net.layers["in"].x = torch.ones(1, 10) + conn.update(reward=1.0, learning=True) + assert torch.allclose( + rule.eligibility_trace[0], torch.tensor([1.0, 0.0, 1.0, 0.0]), atol=1e-6 + ) + + +class TestClampEntersTraces: + """``clamp`` / ``unclamp`` spikes must be reflected in the spike traces.""" + + @staticmethod + def _layer_net(): + net = Network(dt=1.0) + net.add_layer(Input(n=3), "in") + net.add_layer(LIFNodes(n=3, traces=True, tc_trace=20.0), "out") + net.add_connection( + Connection(net.layers["in"], net.layers["out"], w=torch.zeros(3, 3)), + "in", + "out", + ) + return net + + def test_clamped_spike_sets_trace(self): + net = self._layer_net() + mask = torch.tensor([True, False, False]) + net.run(inputs={"in": torch.zeros(1, 3).byte()}, time=1, clamp={"out": mask}) + assert torch.equal(net.layers["out"].s.view(-1), mask) + assert torch.equal(net.layers["out"].x.view(-1), mask.float()) + # And it decays afterwards like any other spike. + net.run(inputs={"in": torch.zeros(1, 3).byte()}, time=1) + assert torch.allclose( + net.layers["out"].x.view(-1), mask.float() * math.exp(-1.0 / 20.0) + ) + + def test_time_indexed_clamp(self): + net = self._layer_net() + mask = torch.zeros(4, 3, dtype=torch.bool) + mask[2, 1] = True + net.run(inputs={"in": torch.zeros(4, 3).byte()}, time=4, clamp={"out": mask}) + x = net.layers["out"].x.view(-1) + assert x[1].item() == pytest.approx(math.exp(-1.0 / 20.0)) + assert x[0].item() == 0.0 and x[2].item() == 0.0 + + def test_unclamped_spike_leaves_no_trace(self): + net = self._layer_net() + net.layers["out"].thresh.fill_(-64.0) + net.run( + inputs={"in": torch.ones(1, 3).byte()}, + time=1, + injects_v={"out": 100.0 * torch.ones(3)}, + unclamp={"out": torch.tensor([False, False, True])}, + ) + s = net.layers["out"].s.view(-1) + assert s[0] and s[1] and not s[2] + assert torch.equal(net.layers["out"].x.view(-1), s.float()) + + +class TestDiehlAndCook2015Rule: + """``DiehlAndCook`` against Diehl & Cook (2015) Sect. 2.3.""" + + @staticmethod + def _reference(pre, post, w0, eta, x_tar, mu, wmax, dt, tc, additive): + w = w0.clone() + x = torch.zeros(pre.shape[1]) + decay = math.exp(-dt / tc) + hist = [] + for t in range(pre.shape[0]): + x = _trace_step(x, pre[t], decay, additive) + # Post-synaptic spikes only: Delta w = eta (x_pre - x_tar) (w_max - w)^mu. + dw = eta * torch.outer(x - x_tar, post[t]) * (wmax - w).clamp(min=0) ** mu + w = (w + dw).clamp(0.0, wmax) + hist.append(w.clone()) + return torch.stack(hist) + + @pytest.mark.parametrize("x_tar", [0.0, 0.3]) + @pytest.mark.parametrize("mu", [1.0, 0.5]) + @pytest.mark.parametrize("additive", [True, False]) + def test_matches_paper_rule(self, x_tar, mu, additive): + dt, tc, wmax, eta = 1.0, 20.0, 1.0, 2e-2 + torch.manual_seed(0) + net = Network(dt=dt) + net.add_layer( + Input(n=12, traces=True, traces_additive=additive, tc_trace=tc), "in" + ) + net.add_layer(LIFNodes(n=6, traces=True, thresh=-60.0), "out") + w0 = 0.5 * torch.rand(12, 6) + conn = Connection( + net.layers["in"], + net.layers["out"], + w=w0.clone(), + nu=(0.0, eta), + update_rule=DiehlAndCook, + wmin=0.0, + wmax=wmax, + x_tar=x_tar, + mu=mu, + ) + net.add_connection(conn, "in", "out") + torch.manual_seed(1) + pre = torch.bernoulli(0.5 * torch.ones(40, 12)) + w_hist, post_hist = [], [] + for t in range(40): + net.run(inputs={"in": pre[t : t + 1].byte()}, time=dt) + post_hist.append(net.layers["out"].s.view(-1).float().clone()) + w_hist.append(conn.w.detach().clone()) + post = torch.stack(post_hist) + assert post.sum() > 0 + w_ref = self._reference(pre, post, w0, eta, x_tar, mu, wmax, dt, tc, additive) + assert (torch.stack(w_hist) - w_ref).abs().max().item() < TOL + + def test_pre_spike_alone_does_not_change_weight(self): + net = Network(dt=1.0) + net.add_layer(Input(n=1, traces=True, traces_additive=True), "in") + net.add_layer(LIFNodes(n=1, traces=True), "out") + conn = Connection( + net.layers["in"], + net.layers["out"], + w=torch.full((1, 1), 0.5), + nu=1.0, + update_rule=DiehlAndCook, + wmin=0.0, + wmax=1.0, + x_tar=0.5, + ) + net.add_connection(conn, "in", "out") + net.layers["out"].x.fill_(0.3) + pre = torch.zeros(3, 1, 1) + pre[1, 0, 0] = 1 + net.run(inputs={"in": pre.byte()}, time=3) + assert conn.w.item() == 0.5 + + def test_x_tar_depresses_silent_inputs(self): + # A post spike with no recent pre spike depresses by eta * x_tar * (wmax - w). + net = Network(dt=1.0) + net.add_layer(Input(n=1, traces=True, traces_additive=True), "in") + net.add_layer(LIFNodes(n=1, traces=True), "out") + conn = Connection( + net.layers["in"], + net.layers["out"], + w=torch.full((1, 1), 0.5), + nu=0.1, + update_rule=DiehlAndCook, + wmin=0.0, + wmax=1.0, + x_tar=0.4, + mu=1.0, + ) + net.add_connection(conn, "in", "out") + net.run( + inputs={"in": torch.zeros(1, 1, 1).byte()}, + time=1, + clamp={"out": torch.tensor([True])}, + ) + assert conn.w.item() == pytest.approx(0.5 - 0.1 * 0.4 * 0.5) + + def test_requires_finite_wmax(self): + net = Network(dt=1.0) + net.add_layer(Input(n=2, traces=True), "in") + net.add_layer(LIFNodes(n=2, traces=True), "out") + with pytest.raises(AssertionError): + Connection( + net.layers["in"], net.layers["out"], nu=0.1, update_rule=DiehlAndCook + ) + + def test_model_opt_in(self): + torch.manual_seed(0) + net = DiehlAndCook2015( + n_inpt=16, + n_neurons=4, + inpt_shape=(1, 4, 4), + learning_rule=MCC_learning.DiehlAndCook, + learning_rule_kwargs={"x_tar": 0.2, "mu": 1.0}, + ) + rule = net.connections[("X", "Ae")].pipeline[0].learning_rule + assert isinstance(rule, MCC_learning.DiehlAndCook) + assert rule.x_tar == 0.2 and rule.mu == 1.0 + assert net.layers["X"].traces_additive # the paper's accumulating trace + w0 = net.connections[("X", "Ae")].pipeline[0].value.clone() + net.run( + inputs={"X": torch.bernoulli(0.5 * torch.ones(30, 1, 1, 4, 4)).byte()}, + time=30, + ) + w = net.connections[("X", "Ae")].pipeline[0].value + # (The model re-normalises each neuron's incoming weights to ``norm`` + # after a run, so only the sign and the change are checked here.) + assert (w >= 0).all() and torch.isfinite(w).all() and not torch.equal(w, w0) + # Default is unchanged: pair STDP with saturating traces. + default = DiehlAndCook2015(n_inpt=16, n_neurons=4, inpt_shape=(1, 4, 4)) + assert isinstance( + default.connections[("X", "Ae")].pipeline[0].learning_rule, + MCC_learning.PostPre, + ) + assert not default.layers["X"].traces_additive diff --git a/test/network/test_mstdp_florian.py b/test/network/test_mstdp_florian.py index c675709de..478a74b76 100644 --- a/test/network/test_mstdp_florian.py +++ b/test/network/test_mstdp_florian.py @@ -11,13 +11,18 @@ both rules against a from-scratch Florian reference (point eligibility, weight update, batch handling, STDP sign, and the ``zero_lag`` timing option). -Reference (i = presynaptic/source, j = postsynaptic/target): - P+_i(t) = P+_i(t-dt) * exp(-dt/tc_plus) + a_plus * pre_i(t) - P-_j(t) = P-_j(t-dt) * exp(-dt/tc_minus) + a_minus * post_j(t) - zeta_ij(t) = P+_i(t) * post_j(t) + pre_i(t) * P-_j(t) - MSTDP: dw_ij(t) = nu * r(t) * zeta_ij(t) - MSTDPET: e_ij(t) = e_ij(t-dt)*exp(-dt/tc_e) + zeta_ij(t-dt)/tc_e - dw_ij(t) = nu * dt * r(t) * e_ij(t) +Reference (i = presynaptic/source, j = postsynaptic/target), Florian (2007) +discrete-time equations; the paper's eq. numbers are given on the right: + P+_i(t) = P+_i(t-dt) * exp(-dt/tc_plus) + a_plus * pre_i(t) (3.11) + P-_j(t) = P-_j(t-dt) * exp(-dt/tc_minus) + a_minus * post_j(t) (3.12) + zeta_ij(t) = P+_i(t) * post_j(t) + pre_i(t) * P-_j(t) (3.10) + MSTDP: w_ij(t+dt) = w_ij(t) + nu * r(t+dt) * zeta_ij(t) (3.9) + MSTDPET: e_ij(t+dt) = e_ij(t)*exp(-dt/tc_e) + zeta_ij(t)/tc_e (2.8) + w_ij(t+dt) = w_ij(t) + nu * dt * r(t+dt) * e_ij(t+dt) (2.7) +i.e. the reward passed at a simulation step multiplies the eligibility built +from the previous step's spikes (``lag=True`` below, the library default). +Defaults tc_plus = tc_minus = 20 ms, tc_e = 25 ms, a_plus = 1, a_minus = -1 +are the paper's Sect. 4.1 values. """ import itertools diff --git a/test/network/test_network.py b/test/network/test_network.py index d2743ef7c..1cfcde255 100644 --- a/test/network/test_network.py +++ b/test/network/test_network.py @@ -66,3 +66,80 @@ def test_add_objects(self, tmp_path): assert ("X", "Y") in _network.connections assert "Y" in _network.monitors del _network + + def test_clone(self): + """ + ``clone()`` round-trips a network in memory. + + Regression test: ``clone()`` called ``torch.load`` without + ``weights_only=False``, so it broke outright when PyTorch 2.6 changed + that default to ``True``. It had no test and no caller, so the breakage + went unnoticed. + """ + import torch + + network = Network(dt=1.0, learning=False) + inpt = Input(10) + network.add_layer(inpt, name="X") + lif = LIFNodes(5) + network.add_layer(lif, name="Y") + w = torch.rand(10, 5) + network.add_connection( + Connection(inpt, lif, w=w.clone()), source="X", target="Y" + ) + + clone = network.clone() + + assert isinstance(clone, Network) + assert clone is not network + assert clone.dt == network.dt + assert clone.learning == network.learning + assert "X" in clone.layers and "Y" in clone.layers + assert ("X", "Y") in clone.connections + assert torch.equal(clone.connections[("X", "Y")].w, w) + + def test_clone_after_save(self, tmp_path): + """ + ``clone()`` still works after ``save()`` has run in the same process. + + Regression test: ``save()`` called + ``torch.serialization.add_safe_globals([self])`` with an instance + instead of a class, which corrupted PyTorch's safe-globals registry and + made every later load in the process fail with + ``'Network' object has no attribute '__qualname__'``. + """ + network = Network(dt=1.0) + network.add_layer(Input(4), name="X") + + network.save(str(tmp_path / "net.pt")) + + clone = network.clone() + assert isinstance(clone, Network) + assert clone.dt == 1.0 + + def test_load_weights_only_parameter(self, tmp_path): + """ + ``load()`` exposes ``weights_only`` and still defaults to ``False``. + + The default has to stay ``False`` because ``save()`` writes the whole + network object, which the safe loader cannot read. ``weights_only=True`` + is offered for files holding plain tensors only, so on a real network + file it is expected to fail rather than silently return something wrong. + """ + import inspect + + signature = inspect.signature(load) + assert "weights_only" in signature.parameters + assert signature.parameters["weights_only"].default is False + + file_path = str(tmp_path / "net.pt") + network = Network(dt=1.0) + network.add_layer(Input(4), name="X") + network.save(file_path) + + # Default path loads a saved network. + assert isinstance(load(file_path), Network) + + # Safe path refuses this file rather than mis-loading it. + with pytest.raises(Exception): + load(file_path, weights_only=True) diff --git a/test/network/test_perf_equivalence.py b/test/network/test_perf_equivalence.py new file mode 100644 index 000000000..0735fed5a --- /dev/null +++ b/test/network/test_perf_equivalence.py @@ -0,0 +1,617 @@ +""" +Regression tests for the performance / memory changes to the per-timestep hot +paths. Every test pins the optimised code to the plain formula it replaced, so a +future change that alters the numerics is caught here. + +Covered: + +* ``MulticompartmentConnection`` can be moved with ``network.to(device)`` + (fixes a ``_apply`` signature mismatch that crashed the Diehl & Cook model). +* ``PostPre`` / ``Hebbian`` on dense ``Connection`` (classic and + multicompartment) use one fused ``addmm_`` instead of materialising the + ``[batch, source.n, target.n]`` outer product. +* ``LearningRule`` skips the ``w *= 1.0`` multiply when no weight decay is set. +* Local-connection rules scale rows directly instead of multiplying by an + ``[n, n]`` identity matrix every timestep. +* Reward-modulated rules cache ``exp(-dt / tc)`` and their default learning + rate tensors. +* Neuron models update their state buffers in place (same operations, same + order) instead of re-assigning module attributes every step. +* ``rank_order`` encoding is vectorised. +""" + +import warnings + +import pytest +import torch + +from bindsnet.encoding import rank_order +from bindsnet.learning import MSTDP, MSTDPET, Hebbian, PostPre, WeightDependentPostPre +from bindsnet.learning import MCC_learning +from bindsnet.learning.learning import ( + _cached_decay, + _dense_outer_update_ok, + _reward_rates, + _row_scale, +) +from bindsnet.models import DiehlAndCook2015 +from bindsnet.network import Network +from bindsnet.network.nodes import ( + AdaptiveLIFNodes, + CurrentLIFNodes, + DiehlAndCookNodes, + Input, + IzhikevichNodes, + LIFNodes, + SRM0Nodes, +) +from bindsnet.network.topology import ( + Connection, + LocalConnection1D, + LocalConnection2D, + LocalConnection3D, + MulticompartmentConnection, +) +from bindsnet.network.topology_features import Weight + + +def _dense_net(rule, batch_size=1, n_in=40, n_out=25, dt=1.0, seed=0, **kwargs): + torch.manual_seed(seed) + net = Network(dt=dt, batch_size=batch_size) + net.add_layer(Input(n=n_in, traces=True), "in") + net.add_layer(LIFNodes(n=n_out, traces=True), "out") + conn = Connection( + net.layers["in"], + net.layers["out"], + nu=kwargs.pop("nu", (1e-2, 2e-2)), + update_rule=rule, + wmin=0.0, + wmax=1.0, + **kwargs, + ) + net.add_connection(conn, "in", "out") + return net, conn + + +def _warm_up(net, steps=15, seed=1): + """Run a few steps so traces and spikes are non-trivial, then return the + snapshot needed by the reference formulas.""" + torch.manual_seed(seed) + b = net.batch_size + n_in = net.layers["in"].n + inp = torch.bernoulli(0.4 * torch.rand(steps, b, n_in)).byte() + net.run(inputs={"in": inp}, time=steps * net.dt) + # Make sure both sides have spikes and traces in the snapshot. + net.layers["in"].s = torch.bernoulli(0.5 * torch.ones(b, n_in)).bool() + net.layers["out"].s = torch.bernoulli( + 0.5 * torch.ones(b, net.layers["out"].n) + ).bool() + net.layers["in"].x = torch.rand(b, n_in) + net.layers["out"].x = torch.rand(b, net.layers["out"].n) + return net + + +def _reference_outer(source, target, batch_size, reduction): + """The un-fused formula: outer products over the batch, then the batch + reduction (``squeeze`` for batch 1, ``sum`` otherwise).""" + source_s = source.s.view(batch_size, -1).unsqueeze(2).float() + source_x = source.x.view(batch_size, -1).unsqueeze(2) + target_s = target.s.view(batch_size, -1).unsqueeze(1).float() + target_x = target.x.view(batch_size, -1).unsqueeze(1) + return ( + reduction(torch.bmm(source_s, target_x), dim=0), # pre: s_pre x x_post + reduction(torch.bmm(source_x, target_s), dim=0), # post: x_pre x s_post + ) + + +def _assert_matches_fused(actual, expected): + """Compare the fused ``addmm_`` result with the un-fused reference formula. + + Both sides call the BLAS library, whose kernel choice depends on the CPU. On + some CPUs the two agree bit for bit; on others (GitHub's CI runners, September + 2026) they differ by float32 rounding. Accept that rounding (the same tolerance + as the batch>1 cases) and report the size of any difference as a warning. + """ + if torch.equal(actual, expected): + return + diff = (actual - expected).abs().max().item() + warnings.warn( + f"fused addmm_ differs from the reference formula by rounding only: " + f"max abs diff {diff:.3e}" + ) + assert torch.allclose(actual, expected, rtol=1e-6, atol=1e-7), diff + + +class TestMulticompartmentDeviceMove: + def test_to_cpu_works(self): + net = DiehlAndCook2015(n_inpt=16, n_neurons=4, inpt_shape=(1, 4, 4)) + net.to("cpu") # crashed before: _apply() got an unexpected ``recurse`` + net.run(inputs={"X": torch.zeros(5, 1, 1, 4, 4).byte()}, time=5) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_to_cuda_moves_features(self): + net = DiehlAndCook2015(n_inpt=16, n_neurons=4, inpt_shape=(1, 4, 4)) + net.to("cuda") + for conn in net.connections.values(): + assert conn.pipeline[0].value.is_cuda + net.run(inputs={"X": torch.zeros(5, 1, 1, 4, 4).byte().cuda()}, time=5) + + +class TestFusedOuterProductRules: + @pytest.mark.parametrize("batch_size", [1, 4]) + def test_postpre_matches_reference(self, batch_size): + net, conn = _dense_net(PostPre, batch_size) + _warm_up(net) + rule = conn.update_rule + w0 = conn.w.detach().clone() + pre, post = _reference_outer( + conn.source, conn.target, batch_size, rule.reduction + ) + expected = (w0 - pre * rule.nu[0] + post * rule.nu[1]).clamp_(0.0, 1.0) + + conn.update(learning=True) + + if batch_size == 1: + _assert_matches_fused(conn.w, expected) + else: + assert torch.allclose(conn.w, expected, rtol=1e-6, atol=1e-7) + + @pytest.mark.parametrize("batch_size", [1, 4]) + def test_hebbian_matches_reference(self, batch_size): + net, conn = _dense_net(Hebbian, batch_size) + _warm_up(net) + rule = conn.update_rule + w0 = conn.w.detach().clone() + pre, post = _reference_outer( + conn.source, conn.target, batch_size, rule.reduction + ) + expected = (w0 + rule.nu[0] * pre + rule.nu[1] * post).clamp_(0.0, 1.0) + + conn.update(learning=True) + + if batch_size == 1: + _assert_matches_fused(conn.w, expected) + else: + assert torch.allclose(conn.w, expected, rtol=1e-6, atol=1e-7) + + def test_fused_path_predicate(self): + # The fused ``addmm_`` path is only taken for dense float32 weights, + # scalar learning rates and the two default batch reductions. + net, conn = _dense_net(PostPre, 1) + rule = conn.update_rule + assert _dense_outer_update_ok(rule, conn.w) + assert not _dense_outer_update_ok(rule, conn.w.to(torch.float64)) + assert not _dense_outer_update_ok(rule, conn.w.to_sparse()) + rule.reduction = torch.mean + assert not _dense_outer_update_ok(rule, conn.w) + rule.reduction = torch.sum + rule.nu = torch.stack([torch.rand(40, 25), torch.rand(40, 25)]) + assert not _dense_outer_update_ok(rule, conn.w) + + def test_custom_reduction_falls_back_and_matches(self): + net, conn = _dense_net(PostPre, 4, reduction=torch.mean) + _warm_up(net) + rule = conn.update_rule + assert rule.reduction is torch.mean + w0 = conn.w.detach().clone() + pre, post = _reference_outer(conn.source, conn.target, 4, torch.mean) + expected = (w0 - pre * rule.nu[0] + post * rule.nu[1]).clamp_(0.0, 1.0) + + conn.update(learning=True) + assert torch.allclose(conn.w, expected, rtol=1e-6, atol=1e-7) + + @pytest.mark.parametrize("dt", [1.0, 0.5]) + def test_mcc_postpre_matches_reference(self, dt): + torch.manual_seed(0) + net = Network(dt=dt) + net.add_layer(Input(n=40, traces=True), "in") + net.add_layer(LIFNodes(n=25, traces=True), "out") + conn = MulticompartmentConnection( + net.layers["in"], + net.layers["out"], + device="cpu", + pipeline=[ + Weight( + "w", + 0.3 * torch.rand(40, 25), + range=[0.0, 1.0], + nu=(1e-2, 2e-2), + learning_rule=MCC_learning.PostPre, + ) + ], + ) + net.add_connection(conn, "in", "out") + _warm_up(net) + rule = conn.pipeline[0].learning_rule + w0 = conn.pipeline[0].value.detach().clone() + pre, post = _reference_outer(conn.source, conn.target, 1, rule.reduction) + # Per-spike increments: no dependence on dt (Morrison et al. 2008 + # eqs. 13-14; the former ``* dt`` factor was removed). + expected = (w0 - pre * rule.nu[0] + post * rule.nu[1]).clamp_(0.0, 1.0) + + conn.update(learning=True) + _assert_matches_fused(conn.pipeline[0].value, expected) + + def test_mcc_hebbian_matches_reference(self): + torch.manual_seed(0) + net = Network(dt=1.0) + net.add_layer(Input(n=40, traces=True), "in") + net.add_layer(LIFNodes(n=25, traces=True), "out") + conn = MulticompartmentConnection( + net.layers["in"], + net.layers["out"], + device="cpu", + pipeline=[ + Weight( + "w", + 0.3 * torch.rand(40, 25), + range=[0.0, 1.0], + nu=(1e-2, 2e-2), + learning_rule=MCC_learning.Hebbian, + ) + ], + ) + net.add_connection(conn, "in", "out") + _warm_up(net) + rule = conn.pipeline[0].learning_rule + w0 = conn.pipeline[0].value.detach().clone() + pre, post = _reference_outer(conn.source, conn.target, 1, rule.reduction) + expected = (w0 + rule.nu[0] * pre + rule.nu[1] * post).clamp_(0.0, 1.0) + + conn.update(learning=True) + _assert_matches_fused(conn.pipeline[0].value, expected) + + +class TestWeightDecay: + def test_no_decay_leaves_weights_untouched(self): + net, conn = _dense_net(PostPre, 1) + assert conn.update_rule.weight_decay == 1.0 + w0 = conn.w.detach().clone() + # No spikes and no traces: the only thing that could change ``w`` is decay. + conn.update(learning=True) + assert torch.equal(conn.w, w0) + + def test_decay_still_applied(self): + net, conn = _dense_net(PostPre, 1, weight_decay=0.01) + w0 = conn.w.detach().clone() + conn.update(learning=True) + assert torch.equal(conn.w, w0 * (1.0 - 0.01)) + + +class TestLocalConnectionRules: + def test_row_scale_equals_diagonal_bmm(self): + torch.manual_seed(0) + vec = torch.rand(3, 17, 1) + mat = torch.rand(3, 17, 9) + reference = torch.bmm(vec * torch.eye(17), mat) + assert torch.equal(_row_scale(vec, mat), reference) + assert torch.equal(_row_scale(vec.squeeze(2), mat), reference) + + @staticmethod + def _local2d_net(rule, batch_size): + torch.manual_seed(0) + net = Network(dt=1.0, batch_size=batch_size) + net.add_layer(Input(shape=[2, 8, 8], traces=True), "in") + net.add_layer(LIFNodes(shape=[3, 3, 3], traces=True), "out") + conn = LocalConnection2D( + net.layers["in"], + net.layers["out"], + kernel_size=4, + stride=2, + n_filters=3, + nu=(1e-2, 2e-2), + update_rule=rule, + wmin=0.0, + wmax=1.0, + ) + net.add_connection(conn, "in", "out") + torch.manual_seed(1) + inp = torch.bernoulli(0.4 * torch.rand(10, batch_size, 2, 8, 8)).byte() + net.run(inputs={"in": inp}, time=10) + net.layers["in"].s = torch.bernoulli( + 0.5 * torch.ones(batch_size, 2, 8, 8) + ).bool() + net.layers["out"].s = torch.bernoulli( + 0.5 * torch.ones(batch_size, 3, 3, 3) + ).bool() + net.layers["in"].x = torch.rand(batch_size, 2, 8, 8) + net.layers["out"].x = torch.rand(batch_size, 3, 3, 3) + return net, conn + + @staticmethod + def _reference_local2d(conn, batch_size): + """The previous implementation: diagonal matrices built with ``torch.eye`` + and multiplied with ``bmm``.""" + kh, kw = conn.kernel_size + sh, sw = conn.stride + c_in = conn.source.shape[0] + n_out = conn.n_filters * conn.conv_size[0] * conn.conv_size[1] + + def unfold(t): + return ( + t.float() + .unfold(-2, kh, sh) + .unfold(-2, kw, sw) + .reshape( + batch_size, conn.conv_size[0] * conn.conv_size[1], c_in * kh * kw + ) + .repeat(1, conn.n_filters, 1) + ) + + target_x = conn.target.x.reshape(batch_size, n_out, 1) * torch.eye(n_out) + target_s = conn.target.s.float().reshape(batch_size, n_out, 1) * torch.eye( + n_out + ) + source_s, source_x = unfold(conn.source.s), unfold(conn.source.x) + red = conn.update_rule.reduction + pre = red(torch.bmm(target_x, source_s), dim=0).view(conn.w.size()) + post = red(torch.bmm(target_s, source_x), dim=0).view(conn.w.size()) + return pre, post + + @pytest.mark.parametrize("batch_size", [1, 2]) + def test_postpre_local2d_matches_diag_reference(self, batch_size): + net, conn = self._local2d_net(PostPre, batch_size) + nu = conn.update_rule.nu + w0 = conn.w.detach().clone() + pre, post = self._reference_local2d(conn, batch_size) + expected = (w0 - nu[0] * pre + nu[1] * post).clamp_(0.0, 1.0) + conn.update(learning=True) + assert torch.equal(conn.w, expected) + + @pytest.mark.parametrize("batch_size", [1, 2]) + def test_hebbian_local2d_matches_diag_reference(self, batch_size): + net, conn = self._local2d_net(Hebbian, batch_size) + nu = conn.update_rule.nu + w0 = conn.w.detach().clone() + pre, post = self._reference_local2d(conn, batch_size) + expected = (w0 + nu[0] * pre + nu[1] * post).clamp_(0.0, 1.0) + conn.update(learning=True) + assert torch.equal(conn.w, expected) + + @pytest.mark.parametrize("batch_size", [1, 2]) + def test_weight_dependent_local2d_matches_diag_reference(self, batch_size): + net, conn = self._local2d_net(WeightDependentPostPre, batch_size) + nu = conn.update_rule.nu + w0 = conn.w.detach().clone() + pre, post = self._reference_local2d(conn, batch_size) + update = -nu[0] * pre * (w0 - conn.wmin) + nu[1] * post * (conn.wmax - w0) + expected = (w0 + update).clamp_(0.0, 1.0) + conn.update(learning=True) + assert torch.equal(conn.w, expected) + + @pytest.mark.parametrize( + "rule", [PostPre, Hebbian, WeightDependentPostPre, MSTDP, MSTDPET] + ) + @pytest.mark.parametrize("dim", [1, 2, 3]) + def test_all_local_rules_run(self, rule, dim): + torch.manual_seed(0) + net = Network(dt=1.0, batch_size=2) + if dim == 1: + in_shape, out_shape = [2, 12], [3, 5] + conn_cls, k = LocalConnection1D, 4 + elif dim == 2: + in_shape, out_shape = [2, 8, 8], [3, 3, 3] + conn_cls, k = LocalConnection2D, 4 + else: + in_shape, out_shape = [2, 6, 6, 6], [2, 2, 2, 2] + conn_cls, k = LocalConnection3D, 4 + net.add_layer(Input(shape=in_shape, traces=True), "in") + net.add_layer(LIFNodes(shape=out_shape, traces=True), "out") + conn = conn_cls( + net.layers["in"], + net.layers["out"], + kernel_size=k, + stride=2, + n_filters=out_shape[0], + nu=(1e-2, 2e-2), + update_rule=rule, + wmin=0.0, + wmax=1.0, + ) + net.add_connection(conn, "in", "out") + inp = torch.bernoulli(0.4 * torch.rand(20, 2, *in_shape)).byte() + kw = {"reward": 0.5} if rule in (MSTDP, MSTDPET) else {} + net.run(inputs={"in": inp}, time=20, **kw) + assert torch.isfinite(conn.w).all() + assert (conn.w >= 0).all() and (conn.w <= 1).all() + if rule in (MSTDP, MSTDPET): + # The traces are kept as vectors now, not diagonal matrices. + r = conn.update_rule + assert r.p_minus.shape == (2, net.layers["out"].n, 1) + assert r.eligibility.shape == (2, *conn.w.shape) + + +class TestRewardRuleCaching: + def test_cached_decay_matches_formula_and_tracks_changes(self): + net, conn = _dense_net(MSTDP, 1, dt=0.5) + rule = conn.update_rule + d = _cached_decay(rule, "tc_plus") + assert torch.equal(d, torch.exp(-0.5 / rule.tc_plus)) + assert _cached_decay(rule, "tc_plus") is d # reused, not recomputed + rule.tc_plus = torch.tensor(7.0) + assert torch.equal( + _cached_decay(rule, "tc_plus"), torch.exp(-0.5 / torch.tensor(7.0)) + ) + conn.dt = 2.0 + assert torch.equal( + _cached_decay(rule, "tc_plus"), torch.exp(-2.0 / torch.tensor(7.0)) + ) + + def test_reward_rates_defaults_reused_and_overrides_honoured(self): + net, conn = _dense_net(MSTDP, 1) + rule = conn.update_rule + ap, am = _reward_rates(rule, {}, torch.device("cpu")) + assert ap.item() == 1.0 and am.item() == -1.0 + ap2, am2 = _reward_rates(rule, {}, torch.device("cpu")) + assert ap2 is ap and am2 is am + ap3, am3 = _reward_rates( + rule, {"a_plus": 0.3, "a_minus": -0.2}, torch.device("cpu") + ) + assert ap3.item() == pytest.approx(0.3) and am3.item() == pytest.approx(-0.2) + + @pytest.mark.parametrize("rule", [MSTDP, MSTDPET]) + def test_reward_rules_match_step_by_step_reference(self, rule): + # From-scratch Florian (2007) reference with the default one-step lag, + # run alongside the network on identical spike trains. + torch.manual_seed(0) + n_in, n_out, T, dt = 12, 7, 25, 0.5 + net = Network(dt=dt) + net.add_layer(Input(n=n_in), "in") + net.add_layer(LIFNodes(n=n_out), "out") + conn = Connection( + net.layers["in"], + net.layers["out"], + nu=1e-2, + update_rule=rule, + wmin=-5.0, + wmax=5.0, + w=0.5 * torch.rand(n_in, n_out), + ) + net.add_connection(conn, "in", "out") + w_ref = conn.w.detach().clone() + r = conn.update_rule + tc_plus, tc_minus, tc_e = float(r.tc_plus), float(r.tc_minus), 25.0 + inp = torch.bernoulli(0.4 * torch.rand(T, n_in)).byte() + p_plus, p_minus = torch.zeros(n_in), torch.zeros(n_out) + elig, e_trace = torch.zeros(n_in, n_out), torch.zeros(n_in, n_out) + reward = 0.7 + for t in range(T): + net.run(inputs={"in": inp[t : t + 1]}, time=dt, reward=reward) + pre = net.layers["in"].s.view(-1).float() + post = net.layers["out"].s.view(-1).float() + if rule is MSTDP: + w_ref += 1e-2 * reward * elig + else: + e_trace = e_trace * torch.exp(torch.tensor(-dt / tc_e)) + elig / tc_e + w_ref += 1e-2 * dt * reward * e_trace + p_plus = p_plus * torch.exp(torch.tensor(-dt / tc_plus)) + pre + p_minus = p_minus * torch.exp(torch.tensor(-dt / tc_minus)) - post + elig = torch.outer(p_plus, post) + torch.outer(pre, p_minus) + w_ref.clamp_(-5.0, 5.0) + assert torch.allclose(conn.w, w_ref, rtol=1e-5, atol=1e-6), t + + +class TestNodesInPlace: + def test_lif_matches_explicit_simulation(self): + torch.manual_seed(0) + n, T, dt = 30, 40, 0.5 + layer = LIFNodes(n=n, traces=True, lbound=-70.0) + net = Network(dt=dt) + net.add_layer(Input(n=n), "in") + net.add_layer(layer, "out") + w = torch.diag(20.0 * torch.ones(n)) + net.add_connection(Connection(net.layers["in"], layer, w=w), "in", "out") + inp = torch.bernoulli(0.6 * torch.rand(T, n)).byte() + + v = layer.rest * torch.ones(1, n) + refrac = torch.zeros(1, n) + x_trace = torch.zeros(1, n) + decay = torch.exp(-torch.tensor(dt) / layer.tc_decay) + trace_decay = torch.exp(-torch.tensor(dt) / layer.tc_trace) + for t in range(T): + net.run(inputs={"in": inp[t : t + 1]}, time=dt) + # Synchronous update: the layer sees the input spikes of step t-1. + prev = inp[t - 1].float() if t > 0 else torch.zeros(n) + inj = (prev @ w).unsqueeze(0) + v = decay * (v - layer.rest) + layer.rest + inj.masked_fill_(refrac > 0, 0.0) + refrac -= torch.tensor(dt) + v += inj + s = v >= layer.thresh + refrac.masked_fill_(s, layer.refrac) + v.masked_fill_(s, layer.reset) + v.masked_fill_(v < -70.0, -70.0) + x_trace *= trace_decay + x_trace.masked_fill_(s, 1.0) + assert torch.equal(layer.s, s), t + assert torch.equal(layer.v, v), t + assert torch.equal(layer.x, x_trace), t + + def test_diehl_and_cook_matches_explicit_simulation(self): + torch.manual_seed(0) + n, T = 20, 40 + layer = DiehlAndCookNodes(n=n, one_spike=False, theta_plus=0.05) + net = Network(dt=1.0) + net.add_layer(Input(n=n), "in") + net.add_layer(layer, "out") + w = torch.diag(15.0 * torch.ones(n)) + net.add_connection(Connection(net.layers["in"], layer, w=w), "in", "out") + inp = torch.bernoulli(0.7 * torch.rand(T, n)).byte() + + v = layer.rest * torch.ones(1, n) + refrac = torch.zeros(1, n) + theta = torch.zeros(n) + decay = torch.exp(-torch.tensor(1.0) / layer.tc_decay) + theta_decay = torch.exp(-torch.tensor(1.0) / layer.tc_theta_decay) + for t in range(T): + net.run(inputs={"in": inp[t : t + 1]}, time=1) + prev = inp[t - 1].float() if t > 0 else torch.zeros(n) + inj = (prev @ w).unsqueeze(0) + v = decay * (v - layer.rest) + layer.rest + theta *= theta_decay + v += (refrac <= 0).float() * inj + refrac -= torch.tensor(1.0) + s = v >= layer.thresh + theta + refrac.masked_fill_(s, layer.refrac) + v.masked_fill_(s, layer.reset) + theta += layer.theta_plus * s.float().sum(0) + assert torch.equal(layer.s, s), t + assert torch.equal(layer.v, v), t + assert torch.equal(layer.theta, theta), t + + @pytest.mark.parametrize( + "node", + [LIFNodes, CurrentLIFNodes, AdaptiveLIFNodes, DiehlAndCookNodes, SRM0Nodes], + ) + def test_state_buffers_keep_identity_and_registration(self, node): + torch.manual_seed(0) + net = Network(dt=1.0) + net.add_layer(Input(n=10), "in") + layer = node(n=6, traces=True) + net.add_layer(layer, "out") + net.add_connection( + Connection(net.layers["in"], layer, w=torch.rand(10, 6)), "in", "out" + ) + v_before = layer.v + net.run(inputs={"in": torch.ones(5, 10).byte()}, time=5) + # In-place updates: same tensor object, still a registered buffer. + assert layer.v is v_before + assert "v" in dict(layer.named_buffers()) + assert torch.isfinite(layer.v).all() + + def test_izhikevich_runs_finite(self): + torch.manual_seed(0) + net = Network(dt=1.0) + net.add_layer(Input(n=10), "in") + layer = IzhikevichNodes(n=10, excitatory=0.8) + net.add_layer(layer, "out") + net.add_connection( + Connection(net.layers["in"], layer, w=5.0 * torch.rand(10, 10)), "in", "out" + ) + net.run( + inputs={"in": torch.bernoulli(0.5 * torch.rand(30, 10)).byte()}, time=30 + ) + assert torch.isfinite(layer.v).all() and torch.isfinite(layer.u).all() + + +class TestRankOrderEncoding: + def test_matches_loop_reference(self): + torch.manual_seed(0) + time = 25 + datum = torch.rand(6, 9) * torch.bernoulli(torch.full((6, 9), 0.7)) + datum[0, 0] = 1.0 # a guaranteed maximum -> earliest spike + out = rank_order(datum.clone(), time=time, dt=1.0) + + # Previous per-neuron loop implementation. + d = datum.flatten().clone() + d /= d.max() + times = torch.zeros(d.numel()) + times[d != 0] = 1 / d[d != 0] + times *= time / times.max() + times = torch.ceil(times).long() + ref = torch.zeros(time, d.numel()).byte() + for i in range(d.numel()): + if 0 < times[i] < time: + ref[times[i] - 1, i] = 1 + assert torch.equal(out, ref.reshape(time, 6, 9)) + assert out.dtype == torch.uint8