STBP is a way to train SNN with datasets by Backward propagation.

Last update: Dec 09, 2022

Overview

Train SNN with STBP in fp32 and low bit(quantize)

Spiking neural network (SNN), compared with depth neural network (DNN), has faster processing speed, lower energy consumption and more biological interpretability, which is expected to approach Strong AI.

STBP is a way to train SNN with datasets by Backward propagation.Using this Repositories allows you to train SNNS with STBP and quantize SNNS with QAT to deploy to neuromorphological chips like Loihi and Tianjic.

Usage

Download via GitHub:

git clone https://github.com/ZLkanyo009/STBP-train-and-compression.git

example to define SNN_layers like ANN_layers

Convert layer to spatiotemporal layer:

conv = nn.Conv2d(...)
conv_s = tdLayer(conv)

Define LIF activation function just like Relu:

spike = LIFSpike()

In the forward function, replace the activation function of each layer with LIF activation function, and replace the calls such as conv() with conv_ s(), then SNN_layers definition is completed.Finally, we use Frequency Coding to decode SNN's output like out = torch.sum(x, dim=2) / steps

def forward(self, x):
    x = self.conv1_s(x)
    x = self.spike(x)
    x = self.pool1_s(x)
    x = self.spike(x)
    x = x.view(x.shape[0], -1, x.shape[4])
    x = self.fc1_s(x)
    x = self.spike(x)
    out = torch.sum(x, dim=2) / steps
    return out

If BN layer is required:

bn = nn.BatchNorm2d(...)
bn = tdBatchNorm(...)
conv_s = tdLayer(conv, bn)

Training Fp32 Model

# Start training fp32 model with: 
# model_name can be ResNet18, CifarNet, ...
python main.py ResNet18 --dataset CIFAR10

# training with DDP:
python -m torch.distributed.launch main.py ResNet18 --local_rank 0 --dataset CIFAR10 --p DDP

# You can manually config the training with: 
python main.py ResNet18 --resume --lr 0.01

Training Quantize Model

# Start training quantize model with: 
# model_name can be ResNet18, CifarNet, ...
python main.py ResNet18 --dataset CIFAR10 -q

# training with DDP:
python -m torch.distributed.launch main.py ResNet18 --local_rank 0 --dataset CIFAR10 -q --p DDP

# You can manually config the training with: 
python main.py ResNet18 -q --resume --bit 4 --lr 0.01

Accuracy

Model	Acc.(fp32)	Acc.(8 bit quantize)
MNISTNet	97.96%	97.57%
ResNet18	84.40%	84.23%

STBP is a way to train SNN with datasets by Backward propagation.

Related tags

Overview

Train SNN with STBP in fp32 and low bit(quantize)

Usage

Download via GitHub:

example to define SNN_layers like ANN_layers

Training Fp32 Model

Training Quantize Model

Accuracy

About STBP

Owner

Ling Zhang

Video Background Music Generation with Controllable Music Transformer (ACM MM 2021 Oral)

This repository contains codes of ICCV2021 paper: SO-Pose: Exploiting Self-Occlusion for Direct 6D Pose Estimation

PRIME: A Few Primitives Can Boost Robustness to Common Corruptions

A general python framework for visual object tracking and video object segmentation, based on PyTorch

Control-Raspberry-Pi-Robot-using-Hand-Gestures - A 4WD Robot car based on Raspberry Pi that controlled by hand gestures(using openCV and mediapipe)

Deep Implicit Moving Least-Squares Functions for 3D Reconstruction

The code repository for EMNLP 2021 paper "Vision Guided Generative Pre-trained Language Models for Multimodal Abstractive Summarization".

High-Fidelity Pluralistic Image Completion with Transformers (ICCV 2021)

A CV toolkit for my papers.

Predict stock movement with Machine Learning and Deep Learning algorithms

Aggragrating Nested Transformer Official Jax Implementation

NeRD: Neural Reflectance Decomposition from Image Collections

Capture all information throughout your model's development in a reproducible way and tie results directly to the model code!

This repo includes our code for evaluating and improving transferability in domain generalization (NeurIPS 2021)

Prefix-Tuning: Optimizing Continuous Prompts for Generation

Enhancing Knowledge Tracing via Adversarial Training

Implementation of Vision Transformer, a simple way to achieve SOTA in vision classification with only a single transformer encoder, in Pytorch

a grammar based feedback fuzzer

Compares various time-series feature sets on computational performance, within-set structure, and between-set relationships.

This repo tries to recognize faces in the dataset you created