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

[ICML 2021] "Graph Contrastive Learning Automated" by Yuning You, Tianlong Chen, Yang Shen, Zhangyang Wang

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

The code is for the paper "A Self-Distillation Embedded Supervised Affinity Attention Model for Few-Shot Segmentation"

Implementation of Pix2Seq in PyTorch

PyTorch implementation of the Value Iteration Networks (VIN) (NIPS '16 best paper)

UltraPose: Synthesizing Dense Pose with 1 Billion Points by Human-body Decoupling 3D Model

Pytorch re-implementation of Paper: SwinTextSpotter: Scene Text Spotting via Better Synergy between Text Detection and Text Recognition (CVPR 2022)

ALIbaba's Collection of Encoder-decoders from MinD (Machine IntelligeNce of Damo) Lab

Pytorch implementation of the paper "Topic Modeling Revisited: A Document Graph-based Neural Network Perspective"

[CVPR2021 Oral] UP-DETR: Unsupervised Pre-training for Object Detection with Transformers

Official code for Score-Based Generative Modeling through Stochastic Differential Equations

MoCoPnet - Deformable 3D Convolution for Video Super-Resolution

Simple tool to combine(merge) onnx models. Simple Network Combine Tool for ONNX.

Official codebase used to develop Vision Transformer, MLP-Mixer, LiT and more.

Neural Ensemble Search for Performant and Calibrated Predictions

Sequence-to-Sequence learning using PyTorch

Code for the paper Progressive Pose Attention for Person Image Generation in CVPR19 (Oral).

Code for our paper Aspect Sentiment Quad Prediction as Paraphrase Generation in EMNLP 2021.

Measuring Coding Challenge Competence With APPS

Rethinking the Importance of Implementation Tricks in Multi-Agent Reinforcement Learning