Skip to content

Repository files navigation

脉冲神经网络 (Spiking Neural Network)

PyTorch License Python

项目简介

本项目实现了一个脉冲神经网络(Spiking Neural Network, SNN)框架,专为GPU加速设计。该框架基于物理空间距离和真实信号传导延迟的 Leaky Integrate-and-Fire (LIF) 神经元模型,融合了以下创新技术:

  • 稀疏张量计算:利用PyTorch稀疏张量(Sparse Tensor)优化内存占用
  • 环形缓冲区:采用Ring Buffer技术高效处理异步延迟传播
  • 动态拓扑演化:支持神经发生(Neurogenesis)与突触发生(Synaptogenesis)的自适应网络生长
  • 自适应稳态阈值:内置Homeostasis动力学以稳定网络发火率
  • 无监督学习:基于突触迹的STDP学习法则,实现自组织学习

项目结构

d:\Github\SNN
├── main.py                 # 主实验脚本:端到端整合示例
├── readme.md              # 项目文档
├── core/
│   ├── __init__.py
│   └── network.py         # 核心网络拓扑模块 (DynamicSpatialSNNGPU)
├── data/
│   ├── __init__.py
│   └── encoder.py         # 脉冲编码器模块 (泊松编码、高斯运动刺激)
├── learning/
│   ├── __init__.py
│   └── stdp.py            # STDP学习法则模块
└── utils/
    ├── __init__.py
    ├── metrics.py         # 监控指标模块 (SNNMonitor)
    └── visualization.py   # 可视化工具模块

系统要求

依赖库

  • PyTorch >= 1.8(建议使用GPU版本以获得最佳性能)
  • NumPy >= 1.18
  • Matplotlib >= 3.0
  • CUDA >= 10.2(可选,用于GPU加速)

运行环境

  • 操作系统:Windows / Linux / macOS
  • Python 版本:3.8 及以上
  • GPU显存:建议8GB以上(可选)

快速开始

安装

# 克隆项目
git clone <repository-url>
cd SNN

# 安装依赖(推荐使用虚拟环境)
pip install torch numpy matplotlib

基础使用示例

import torch
from core.network import DynamicSpatialSNNGPU
from learning.stdp import TraceBasedSTDP
from data.encoder import moving_gaussian_stimulus
from utils.metrics import SNNMonitor

# 初始化
device = 'cuda' if torch.cuda.is_available() else 'cpu'
snn = DynamicSpatialSNNGPU(max_delay=10, device=device)

# 神经发生:创建1000个神经元
snn.neurogenesis(num_neurons=1000)

# 突触发生:建立稀疏连接
snn.synaptogenesis(p_max=0.5)

# 时间步进
for t in range(1000):
    spikes = snn.step_evolution(external_current=input_current)

运行完整实验

python main.py

核心模块详细说明

1. 核心网络模块 (core/network.py)

类:DynamicSpatialSNNGPU

基于GPU加速的动态空间脉冲神经网络类。

关键特性:

  • 在三维物理空间中生成并管理神经元
  • 基于欧氏距离和指数衰减概率建立突触
  • 根据物理距离预计算突触延迟
  • 支持异步延迟脉冲传播
  • 内置自适应稳态阈值机制

核心参数:

参数 类型 说明
max_delay int 最大突触传导延迟步数,用于初始化环形缓冲区
decay_lambda float 空间衰减常数 $\lambda$,控制距离对连接概率的影响
signal_velocity float 动作电位传导速度 $v$,与距离共同确定延迟 $\Delta t = \lceil d/v \rceil$
device str 计算设备,'cuda'或'cpu'

全局状态张量:

张量 形状 说明
positions $(N, 3)$ 神经元的三维物理坐标
v $(N,)$ 膜电位状态向量
v_th $(N,)$ 动态发射阈值向量(Homeostasis)
spike_buffer $(max_delay, N)$ 环形缓冲区,记录历史脉冲
synapse_weights $(N, N)$ 稀疏突触权重张量
edge_indices $(2, E)$ 突触连接索引(源、目标)
delay_indices $(E,)$ 每条突触的具体延迟步数

核心方法:

neurogenesis(num_new_neurons, space_bounds)
  • 功能:在三维空间中批量生成新神经元
  • 参数
    • num_new_neurons (int):生成的神经元数量
    • space_bounds (tuple):三维空间边界,默认 $(100, 100, 100)$
  • 机制:使用 torch.rand 生成随机坐标,通过 torch.cat 动态拼接到状态张量
synaptogenesis(p_max, batch_size)
  • 功能:根据空间距离建立突触连接
  • 参数
    • p_max (float):距离为0时的最大连接概率
    • batch_size (int):分块计算的批大小,防止显存溢出
  • 连接规则:$P(i \to j) = P_{max} \exp(-D_{ij}/\lambda)$
  • 延迟计算:$\Delta t_{ij} = \lceil D_{ij}/v \rceil$
step_evolution(external_current)
  • 功能:执行一个时间步的网络演化
  • 参数external_current (Tensor):外部注入电流,形状 $(N,)$
  • 返回值:当前时间步的脉冲活动,形状 $(N,)$ 的布尔张量

LIF神经元动力学:

$$V(t+1) = \beta \cdot (V(t) - V_{rest}) + V_{rest} + I_{syn}(t) + I_{ext}(t)$$

其中:

  • $\beta = \exp(-1/\tau_m)$ 为膜时间常数衰减因子
  • $I_{syn}(t) = \sum_{j} W_{ij} \cdot S_j(t - \Delta t_{ij})$ 为突触电流
  • $I_{ext}(t)$ 为外部注入电流

自适应阈值(Homeostasis)机制:

$$\theta(t+1) = \theta_{base} + \beta_{th} \cdot (\theta(t) - \theta_{base})$$

发放时阈值暴涨:

$$\theta[\text{spikes}] \mathrel{+}= \theta_{plus}$$


2. 数据编码模块 (data/encoder.py)

函数:poisson_encode(stimulus, time_steps, dt)

  • 功能:将连续信号编码为泊松脉冲序列
  • 参数
    • stimulus (Tensor):预期发放率,形状 $(N,)$
    • time_steps (int):模拟时间步
    • dt (float):时间步长(毫秒)
  • 返回值:形状 $(time_steps, N)$ 的布尔脉冲张量

函数:moving_gaussian_stimulus(positions, time_steps, start_pos, end_pos, max_rate, sigma, dt)

  • 功能:生成三维空间中运动的高斯刺激
  • 参数
    • positions (Tensor):神经元坐标,形状 $(N, 3)$
    • start_pos, end_pos (tuple):刺激运动的起终点
    • max_rate (float):刺激中心的最高发放率(0-1)
    • sigma (float):高斯感受野的空间衰减半径
  • 输出:时空强相关的泊松脉冲序列,极利于STDP学习

高斯函数:

$$P(i,t) = P_{max} \exp\left(-\frac{d_{i,center(t)}^2}{2\sigma^2}\right)$$


3. 学习模块 (learning/stdp.py)

类:TraceBasedSTDP

基于突触迹的脉冲时序依赖可塑性(Spike-Timing-Dependent Plasticity, STDP)学习器。

关键特性:

  • 维护所有神经元的前/后神经突触迹
  • 完全基于稀疏拓扑的高效权重更新
  • 支持LTP与LTD的非对称学习
  • GPU并行化计算

核心参数:

参数 类型 说明
tau_plus float LTP时间衰减常数 $\tau^+$
tau_minus float LTD时间衰减常数 $\tau^-$
a_plus float LTP学习率与权重增加上限 $A^+$
a_minus float LTD学习率与权重减少上限 $A^-$
w_max float 突触权重硬上限
w_min float 突触权重硬下限

突触迹动力学:

$$\mathrm{trace}^{pre}(t+1) = \mu^+ \cdot \mathrm{trace}^{pre}(t) + S_{pre}(t)$$

$$\mathrm{trace}^{post}(t+1) = \mu^- \cdot \mathrm{trace}^{post}(t) + S_{post}(t)$$

其中 $\mu^{\pm} = \exp(-1/\tau^{\pm})$

权重更新规则:

  • LTP阶段(突触后神经元发放):$\Delta W \mathrel{+}= A^+ \cdot \mathrm{trace}^{pre} \cdot S_{post}$
  • LTD阶段(突触前神经元发放):$\Delta W \mathrel{-}= A^- \cdot \mathrm{trace}^{post} \cdot S_{pre}$

核心方法:

update_traces(spikes)
  • 更新所有神经元的前/后突触迹
  • 应在每个时间步后立即调用
update_weights(synapse_weights, edge_indices, spikes)
  • 执行STDP权重更新
  • 支持稀疏张量高效计算

4. 监控模块 (utils/metrics.py)

类:SNNMonitor

SNN训练过程的监控与分析工具。

关键特性:

  • 记录时间序列的平均发放率
  • 保存权重分布的时间快照
  • 生成学术规范的图表并保存

核心方法:

record_firing_rate(spikes)
  • 记录当前时间步全网的平均发放率
  • 参数:spikes (Tensor):形状 $(N,)$ 的脉冲布尔张量
record_weights(step, weights)
  • 记录指定时间步的突触权重快照
  • 参数:
    • step (int):时间步编号
    • weights (Tensor):支持稠密张量或稀疏张量
plot_firing_rate(window_size, title, save_path)
  • 绘制网络平均发放率随时间的演化曲线
  • 支持滑动窗口平滑
plot_weight_distribution(title, save_path)
  • 绘制权重分布的直方图
  • 展示不同时间步的权重极化演化

5. 可视化模块 (utils/visualization.py)

函数:plot_raster(spike_history, title)

  • 功能:绘制脉冲光栅图(Raster Plot)
  • 参数
    • spike_history (Tensor):形状 $(time_steps, num_neurons)$ 的脉冲布尔张量
    • title (str):图像标题
  • 用途:可视化神经网络的脉冲放电时空模式

函数:plot_3d_topology(positions, edge_indices, weights, max_edges, title)

  • 功能:在三维空间中可视化神经元节点和突触连接
  • 参数
    • positions (Tensor):神经元坐标,形状 $(N, 3)$
    • edge_indices (Tensor):连接索引,形状 $(2, E)$
    • weights (Tensor):突触权重,用于映射连接透明度
    • max_edges (int):限制绘制的最大连线数,防止视觉过载
  • 特点
    • 神经元节点使用蓝色标记
    • 连线透明度根据权重归一化映射
    • 支持交互式3D旋转

类:OfflineTopologyAnimator

离线SNN拓扑演化动画生成器。

功能特性:

  • 轻量级保存模拟过程中的网络状态快照
  • 批量渲染生成高质量3D旋转视频
  • 利用Line3DCollection优化大规模突触的绘制速度
  • 权重动态映射到线的宽度和颜色

核心方法:

save_checkpoint(time_step, positions, edge_indices, weights)
  • 保存当前时间步的网络拓扑快照
  • 建议每50-200步调用一次以平衡内存与平滑度
generate_animation(output_filename, max_edges, title)
  • 从已保存的快照渲染生成MP4或GIF动画
  • 参数:
    • output_filename (str):输出文件名(支持.mp4或.gif)
    • max_edges (int):每帧绘制的最大突触数(按权重排序)
    • title (str):视频标题

依赖要求: 系统需要安装 FFmpeg 用于视频编码


使用示例

示例1:基础网络演化

import torch
from core.network import DynamicSpatialSNNGPU

# 初始化网络
snn = DynamicSpatialSNNGPU(max_delay=10, device='cuda')

# 创建1000个神经元
snn.neurogenesis(1000)

# 建立突触连接
snn.synaptogenesis(p_max=0.5)

# 单步演化
external_current = torch.randn(1000, device='cuda') * 0.1
spikes = snn.step_evolution(external_current=external_current)

print(f"当前时间步发放率: {spikes.float().mean().item():.3f}")

示例2:STDP学习

from learning.stdp import TraceBasedSTDP

# 初始化STDP学习器
stdp = TraceBasedSTDP(num_neurons=1000, device='cuda')

# 学习过程
for t in range(1000):
    spikes = snn.step_evolution()
    stdp.update_traces(spikes)
    new_weights = stdp.update_weights(snn.synapse_weights, snn.edge_indices, spikes)
    snn.synapse_weights = new_weights

示例3:完整实验(见 main.py

python main.py

该脚本执行以下操作:

  1. 初始化1000神经元的网络
  2. 生成空间运动的高斯刺激
  3. 运行10000步STDP自组织学习
  4. 生成统计图表(发放率、权重分布)

核心概念快速参考

LIF神经元模型

  • 漏积分发火:膜电位随时间指数衰减,脉冲时电位重置
  • 异步乱序处理:通过环形缓冲区支持可变的突触延迟

STDP学习规则

  • 时间因果性:突触前脉冲领先后脉冲 → LTP增强
  • 时间反因果性:突触前脉冲滞后后脉冲 → LTD减弱
  • 突触迹机制:维护神经元的"发放历史"用于权重更新

Homeostasis稳态

  • 维持网络稳定的平均发放率
  • 防止神经元过度激活或沉默
  • 通过自适应阈值实现

性能优化建议

优化 方法
显存峰值 使用稀疏张量 + 分块突触发生 + 降频权重记录
计算速度 GPU并行化 + 向量化操作 + 避免CPU-GPU转移
可视化 降低绘图频率 + 限制绘制的边数量 + 异步保存

常见问题 (FAQ)

Q: 为什么突触发生后没有连接?
A: 检查 p_max 参数。如果离得太远,指数衰减会导致连接概率近似为0。

Q: 如何加速生成动画?
A: 降低 checkpoint_interval 的频率,或减少 max_edges 参数。

Q: STDP学习不收敛?
A: 调整 a_plus, a_minus, tau_plus, tau_minus 等参数,或增加学习时间。


许可证

MIT License - 详见 LICENSE 文件

联系方式

如有问题或建议,欢迎提交Issue或Pull Request。

About

此项目实现了一个简易的脉冲神经网络(SNN)框架

Resources

Stars

3 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages