并行脑仿真#

NeurAI提供神经网络的分布式训练及大规模脑仿真并行功能，用户只需要通过 neurai.parallelize.parallelize 装饰器完成所有分布式相关功能的调用。具体功能参见这篇分布式训练。

本文将介绍如何使用 neurai.parallelize.parallelize 来实现并行脑仿真。

并行装饰器介绍#

其原理是使用 @parallelize 装饰器，在装饰器中增加了代理神经元和远程突触，实现了虚拟创建，来实现并行脑仿真。neurai.parallelize.parallelize 是一个装饰器，在装饰器内部把 neurai.nn.neuron.exp_lif.ExpLIF 对象，自动替换为代理类对象 neurai.parallelize.neuron.exp_lif.ProxyExpLIF，以实现与 `NeurAI` 中非并行脑仿真方法的接口统一。

搭建脑仿真网络#

要在 NeurAI 中运行并行脑仿真，首先导入模块 neurai.parallelize。然后按顺序构建神经元和突触模型，可以按以下方式设置它：

from neurai.parallelize import parallelize
from neurai.nn.neuron import ExpLIF, PoissonGenerator
from neurai.nn.synapse import StaticSynapse
from neurai import nn
from neurai.monitor import MonitorBS, MonitorConf, MemoryRecorder
from neurai.util import visualization
import matplotlib.pyplot as plt
import jax.numpy as jnp


size = 10
class SimpleNet(nn.SNet):

  @parallelize
  def setup(self):
    # 创建ExpLIF神经元簇
    self.lif0 = ExpLIF(size=size, v_init=-55, V_rest=-60.0, V_th=-50.0, V_reset=-60.0, tau=20.0, R=1, I_e=200.0,)
    self.lif1 = ExpLIF(size=size, v_init=-55, V_rest=-60.0, V_th=-51.0, V_reset=-60.0, tau=20.0, R=1, I_e=100.0,)
    self.lif2 = ExpLIF(size=size, v_init=-55, V_rest=-60.0, V_th=-52.0, V_reset=-60.0, tau=20.0, R=1, I_e=50.0,)
    self.lif3 = ExpLIF(size=size, v_init=-55, V_rest=-60.0, V_th=-53.0, V_reset=-60.0, tau=20.0, R=1, I_e=0.0,)
    self.lif4 = ExpLIF(size=size, v_init=-55, V_rest=-60.0, V_th=-54.0, V_reset=-60.0, tau=20.0, R=1, I_e=0.0,)
    self.lif5 = ExpLIF(size=size, v_init=-55, V_rest=-60.0, V_th=-55.0, V_reset=-60.0, tau=20.0, R=1, I_e=0.0,)
    # 创建泊松生成器
    self.p = PoissonGenerator(size=size, rate=10000)

    # 创建突触模型
    self.synapse1 = StaticSynapse(pre=self.lif3, post=self.lif0, conn=nn.One2One(), weight=10.0, delay_step=2,)
    self.synapse2 = StaticSynapse(pre=self.lif1, post=self.lif4, conn=nn.One2One(), weight=10.0, delay_step=2,)
    self.synapse3 = StaticSynapse(pre=self.lif2, post=self.lif5, conn=nn.One2One(), weight=500.0, delay_step=2,)

    self.conn1 = StaticSynapse(self.p, self.lif0, conn=nn.One2One(), weight=10, delay_step=2)

class AllModule(nn.Module):

  def setup(self):
    self.network = nn.SNetLayer(SimpleNet)

  def __call__(self, input=None, t=0, monitor=None):
    _, mon = self.network(input, t=t, monitor=monitor)

    return mon

设置参数和监视器#

首先设置仿真的总持续时间，再设置监视器。

sim_t = 500.0
monitorbs = MonitorBS(
  monitors=[
    MonitorConf("ExpLIF_0", "spike"), MonitorConf("ExpLIF_0", "v"),
    MonitorConf("ExpLIF_1", "spike"), MonitorConf("ExpLIF_1", "v"),
    MonitorConf("ExpLIF_2", "spike"), MonitorConf("ExpLIF_2", "v"),
    MonitorConf("ExpLIF_3", "spike"), MonitorConf("ExpLIF_3", "v"),
    MonitorConf("ExpLIF_4", "spike"), MonitorConf("ExpLIF_4", "v"),
    MonitorConf("ExpLIF_5", "spike"), MonitorConf("ExpLIF_5", "v"),
  ],
  recorder=MemoryRecorder(),
)

有关监视器的详细信息，可参考监视神经元及突触信息。

运行脑仿真网络#

通过 init 传递初始参数并完成网络初始化。使用 neurai.nn.module.run 运行网络仿真，并输出监视参数。在 neurai.parallelize 中，没有合并监视参数，输出可以通过可视化或在每个进程中存储监视参数来查看。

# 实例化网络并初始化参数
net = AllModule()
param = net.init()
# 运行网络仿真
mon = net.run(param, t=sim_t, monitor=monitorbs)

fig = plt.figure(figsize=(4, 12), dpi=150)

for i, key in enumerate(mon.keys()):
  if key == 'ts':
    continue
  ax = fig.add_subplot(len(mon) - 1, 1, i + 1)
  if key[-1] == 'v':
    visualization.line_plot(mon["ts"], jnp.asarray(mon[key])[:, 0:size], ax=ax, xlabel="Time(ms)",
                            ylabel=key + "(mv)", show=False, save=False, title=key,)
  elif key[-1] == 'e':
    visualization.raster_plot(mon["ts"], jnp.asarray(mon[key]), ax=ax, show=False, save=False, title=key)

在以上设置都完成后，可以通过执行 python3 parallelize_example.py 来单进程运行，也可以通过 mpirun 来多进程仿真：

mpirun -np ${nums} python3 parallelize_example.py

并行脑仿真

Contents

并行脑仿真#

并行装饰器介绍#

搭建脑仿真网络#

设置参数和监视器#

运行脑仿真网络#