文章介绍了BrainPy这个计算神经科学建模工具，强调了其JIT编译加速功能，通过对比展示了JIT如何提高代码执行效率。此外，还提到了BrainPy的Runner类自动化JIT编译以及基础数据操作，包括数组和动态变量的使用，并给出了控制流的实现方法，如where和for/while_loop语句。
参考书目：《神经计算建模实战——基于brainpy》 吴思
brainpy是一套基于脑动力学模型定义、训练和分析的通用编程软件，致力于为计算神经建模提供新生态。
1 brainpy的安装
BrainPy下载链接
https://brainpy.readthedocs.io/en/latest/quickstart/installation.html
一定要按照官网上的步骤，否则安装真的好容易出bug。安装完成后，之后的代码均需要引入相关环境：
import brainpy as bp
import brainpy.math as bm
import numpy as np
import jax.numpy as jnp
import time
import matplotlib.pyplot as plt
bm.set_platform('cpu') #利用cpu进行计算2 JIT编译加速
JIT基础知识
BrainPy的核心思想之一是即时编译（Just-In-Time compilation），JIT编译可以视为一种加速机制。在使用JIT之前，需要将目标函数或类用brainpy.math.jit()进行包装。以GELU函数的执行为例，测试在JIT和非JIT编译下的运行速度：
def gelu(x): # 高斯误差线性单元函数
sqrt = bm.sqrt(2 / bm.pi)
cdf = 0.5 * (1.0 + bm.tanh( sqrt * (x+0.044715 * (x ** 3))))
y = x*cdf
return y
x=bm.random.random(100000)
sta=time.clock()
gelu(x)
end=time.clock()
print("运行耗时", end-sta)
gelu_jit = bm.jit(gelu)
sta=time.clock()
gelu_jit(x) #使用JIT
end=time.clock()
print("运行耗时", end-sta)得到的输出为：
运行耗时 0.1279861999999996
运行耗时 0.032745799999999825由此可见，JIT编译会大大提升执行速度。上面的代码也展示了JIT的用法。但是在一个动力系统中，众多的动态变量和微分方程会使得计算变得十分复杂，因此，BrainPy允许对类对象进行JIT编译。其中，动态变量必须被定义为brainpy.math.Variable，采用Logistic回归分类器（其中权重w需要定义为动态变量）作为例子，并在不使用/使用JIT编译下测试运行时间：
# 类继承于brainpy.Base
class LogisticRegression(bp.Base):
def __init__(self,dimension):
super(LogisticRegression,self).__init__()
self.dimension = dimension #参数定义
self.w = bm.Variable(2.0 * bm.ones(dimension) - 1.3) #动态变量定义
def __call__(self,X,Y):
u = bm.dot (( (1.0 / (1.0 + bm.exp(-Y * bm.dot(X , self.w)))-1.0)* Y),X)
self.w.value = self.w - u #更新动态变量
#计算执行时间的函数benckmark
def benckmark(model , point , labels , num_iter = 30, name=''):
t0=time.time()
for i in range(num_iter):
model(points , labels)
print(f' {name} used time {time.time() - t0} s \n' )
#定义数据集和标签
num_dim,num_points=10,20000000
points = bm.random.random((num_points , num_dim))
labels = bm.random.random(num_points)
#不使用JIT
lr1 = LogisticRegression(num_dim)
benckmark(lr1 , points ,labels,name='without jit')
#使用JIT
lr2 = LogisticRegression(num_dim)
lr2=bm.jit((lr2))
benckmark(lr2 , points ,labels,name='with jit'JIT编译的使用方式类似于函数，只需将类实例传入jit()中。运行结果如下：
without jit used time 11.542979955673218 s
with jit used time 6.864618301391602 s 事实上，针对规模较大的模型，JIT编译带来的加速通常要显著得多。
brainpy.Runner类
实际的大规模动力学系统常常包含大量神经元和突触模型，如果显式地将每个对象都包装到brainpy.math.jit()中，则会显得十分繁琐。为了简化编程逻辑，BrainPy实现了自动JIT编译的技术。brainpy.Runner类是模拟、训练、积分等各种运行器的基类，只要将目标工程传入Runner中，Runner就可以自动JIT编译目标工程。以HH神经元模型（3.1-Hodgkin-Huxley (HH) 模型 - 知乎 (zhihu.com)
https://zhuanlan.zhihu.com/p/102770891）为例：
#HH模型 Runner
model = bp.neurons.HH(100) #100个神经元
runner = bp.DSRunner(target=model,inputs=('input',10.)) #jit默认为True
print(runner(duration=100,eval_time=True)[0]) #模拟100ms
runner = bp.DSRunner(target=model,inputs=('input',10.),jit=False) #关闭GIT
print(runner(duration=100,eval_time=True)[0]) 输出结果为：
Predict 1000 steps: : 100%|██████████| 1000/1000 [00:00<00:00, 2074.83it/s]
0.4809699058532715
Predict 1000 steps: : 100%|██████████| 1000/1000 [02:06<00:00, 7.93it/s]
126.1071035861969可见不使用JIT的话，运行速度真的很慢。
3 基础的数据操作
除了python中的基本数据类型，BrainPy中有两类特殊的数据类型：数组（Array）和动态变量（Variable）
数组（Array）
数组是一种数据结构。在brainpy中，该数据结构是一个包含相同数据类型的多维数组，常见数值或布尔类型。可用brainpy.math.array()创建一个一维或多维数组并查看数组本身的几种属性：
- .ndim：维度
- .shape：数组形状 例如n行m列矩阵的形状为(n,m)
- .size：数组的元素总数，等于shape中各元素乘积
- .dtype：描述数组中元素类型
使用实例如下：
t1 = bm.array([ [[0,1,2,3],[1,2,3,4],[4,5,6,7]],
[[3,2,4,5],[23,5,8,9],[2,-2,-2,-3]]])
print(t1)
>>>Array(value=DeviceArray([[[ 0, 1, 2, 3],
[ 1, 2, 3, 4],
[ 4, 5, 6, 7]],
[[ 3, 2, 4, 5],
[23, 5, 8, 9],
[ 2, -2, -2, -3]]]),
dtype=int32)
print('t1.ndim: {}'.format(t1.ndim))
>>>t1.ndim: 3
print('t1.shape: {}'.format(t1.shape))
>>>t1.shape: (2, 3, 4)
print('t1.size: {}'.format(t1.size))
>>>t1.size: 24
print('t1.dtype: {}'.format(t1.dtype))
>>>t1.dtype: int32
动态变量（Variable）
数组的问题在于一旦被交给JIT编译器，数组里面的值便不能修改。而动态变量中的数据可以被修改。可以用brainpy.math.Variable(t)将数组t转换为动态变量。动态变量的就地更新方法如下：
v=bm.Variable(bm.arange(4))
print(v)
>>>Variable(value=DeviceArray([0, 1, 2, 3]), dtype=int32)
v[0]=10 #索引
print(v)
>>>Variable(value=DeviceArray([10, 1, 2, 3]), dtype=int32)
v[1:3]=8 #切片
print(v)
>>>Variable(value=DeviceArray([10, 8, 8, 3]), dtype=int32)
v**=2 #增量赋值
print(v)
>>>Variable(value=DeviceArray([100, 64, 64, 9]), dtype=int32)
v.value = bm.arange(4) #.value赋值
print(v)
>>>Variable(value=DeviceArray([0, 1, 2, 3]), dtype=int32)
v.update(bm.array([2,34,1,9])) #.update赋值
print(v)
>>>Variable(value=DeviceArray([ 2, 34, 1, 9]), dtype=int32)4 控制流
条件语句
如果不依赖于动态变量，可正常编写控制流。但当判断条件是动态变量时，python通用的if/else语句便不再适用。可用如下方法代替：
- brainpy.math.where(condition,x,y)：条件真返回x，否则返回y
- brainpy.math.ifelse(conditions,branches,operands,dyn_vars=None) ：条件语句、分支语句、分支语句需要的参数、分支语句中出现的动态变量
循环语句
- brainpy.math.for_loop()可用来实现for循环语句：
运行实例如下：
class LoopStruct(bp.Base):
def __init__(self):
super(LoopStruct,self).__init__()
rng = bm.random.RandomState(123)
self.seq = rng.random(1000)
self.res = bm.Variable(bm.zeros(1))
def __call__(self):
def add(s):
self.res +=s
bm.for_loop(add,dyn_vars=[self.res],operands=self.seq)
return self.res.value
model = bm.jit(LoopStruct())
print(measure_time(model)) #第一次运行会触发编译
>>>Result: [501.74664] , Time:0.03390955924987793
print(measure_time(model)) #第二次运行不会触发编译
>>>Result: [1003.4931] , Time:0.02.brainpy.math.while_loop()语句实现while功能，实例如下：
a = bm.Variable(bm.zeros(1))
b = bm.Variable(bm.ones(1))
def cond(x,y):
return x<6.
def body(x,y):
a.value += x
b.value *= y
return x + b[0],y +1.
res = bm.while_loop(body , cond ,dyn_vars=[a,b],operands=(1.,1.))
print(res)
>>>(DeviceArray(10., dtype=float32), DeviceArray(4., dtype=float32, weak_type=True))
print(a)
>>>Variable(value=DeviceArray([7.]), dtype=float32)
print(b)
>>>Variable(value=DeviceArray([6.]), dtype=float32)该程序一共进行了3次循环迭代，每次循环后执行结果如下：
| 初始值 | 第一次循环 | 第二次循环 | 第三次循环 |
| a=0 | a=a+x=1 | a=a+x=3 | a=a+x=7 |
| b=1 | b=b*y=1 | b=b*y=2 | b=b*y=6 |
| res | x=1 y=1 | x=x+b=2 y=y+1=2 | x=x+b=4 y=y+1=3 |
| cond | True | True | False (a>6) |
扫一扫
5581
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
