update examples

Author: chaoming0625

examples/dynamics_simulation/Wang_2002_decision_making_spiking.py

@@ -267,11 +267,11 @@ def batching_run():
     runner = bp.DSRunner(
       net,
       monitors=['A.spike', 'B.spike', 'IA.freq', 'IB.freq'],
-      time_major=False
+      data_first_axis=False
     )
     runner.run(total_period)
 
-  coherence = coherence.to_numpy()
+  coherence = bm.as_numpy(coherence)
   fig, gs = bp.visualize.get_figure(num_row, num_col, 3, 4)
   for i in range(num_row):
     for j in range(num_col):

examples/training_snn_models/SurrogateGrad_lif_fashion_mnist.py

@@ -163,7 +163,8 @@ def loss_fun(predicts, targets):
     return loss + l2_loss + l1_loss
 
   trainer = bp.train.BPTT(
-    model, loss_fun,
+    model,
+    loss_fun,
     optimizer=bp.optim.Adam(lr=lr),
     monitors={'r.spike': net.r.spike},
   )

examples/training_snn_models/fashion_mnist_conv_lif.py

@@ -11,6 +11,7 @@
 
 import brainpy as bp
 import brainpy.math as bm
+from brainpy.tools import DotDict
 
 bm.set_environment(mode=bm.training_mode, dt=1.)
 
@@ -82,7 +83,8 @@ def reset_state(self, batch_size):
   def update(self, s, x):
     self.V.value += x
     spike = self.spike_fun(self.V - self.v_threshold)
-    s = lax.stop_gradient(spike)
+    # s = lax.stop_gradient(spike)
+    s = spike
     if self.reset_mode == 'hard':
       one = lax.convert_element_type(1., bm.float_)
       self.V.value = self.v_reset * s + (one - s) * self.V
@@ -97,24 +99,24 @@ def __init__(self, n_time: int, n_channel: int):
     self.n_time = n_time
 
     self.block1 = bp.Sequential(
-      bp.layers.Conv2D(1, n_channel, kernel_size=3, padding=(1, 1), b_initializer=None),
+      bp.layers.Conv2D(1, n_channel, kernel_size=3, padding=(1, 1), ),
       bp.layers.BatchNorm2D(n_channel, momentum=0.9),
       IFNode((28, 28, n_channel), spike_fun=bm.surrogate.arctan)
     )
     self.block2 = bp.Sequential(
       bp.layers.MaxPool([2, 2], 2, channel_axis=-1),  # 14 * 14
-      bp.layers.Conv2D(n_channel, n_channel, kernel_size=3, padding=(1, 1), b_initializer=None),
+      bp.layers.Conv2D(n_channel, n_channel, kernel_size=3, padding=(1, 1), ),
       bp.layers.BatchNorm2D(n_channel, momentum=0.9),
       IFNode((14, 14, n_channel), spike_fun=bm.surrogate.arctan),
     )
     self.block3 = bp.Sequential(
       bp.layers.MaxPool([2, 2], 2, channel_axis=-1),  # 7 * 7
       bp.layers.Flatten(),
-      bp.layers.Dense(n_channel * 7 * 7, n_channel * 4 * 4, b_initializer=None),
+      bp.layers.Dense(n_channel * 7 * 7, n_channel * 4 * 4,),
       IFNode((4 * 4 * n_channel,), spike_fun=bm.surrogate.arctan),
     )
     self.block4 = bp.Sequential(
-      bp.layers.Dense(n_channel * 4 * 4, 10, b_initializer=None),
+      bp.layers.Dense(n_channel * 4 * 4, 10, ),
       IFNode((10,), spike_fun=bm.surrogate.arctan),
     )
 
@@ -138,8 +140,6 @@ def main():
   parser.add_argument('-data-dir', default='./data', type=str, help='root dir of Fashion-MNIST dataset')
   parser.add_argument('-out-dir', default='./logs', type=str, help='root dir for saving logs and checkpoint')
   parser.add_argument('-lr', default=0.1, type=float, help='learning rate')
-  parser.add_argument('-save-es', default=None,
-                      help='filepath for saving a batch spikes encoded by the first {Conv2d-BatchNorm2d-IFNode}')
   args = parser.parse_args()
   print(args)
 
@@ -163,20 +163,20 @@ def main():
   def inference_fun(X, fit=True):
     net.reset_state(X.shape[0])
     return bm.for_loop(lambda sha: net(sha.update(dt=bm.dt, fit=fit), X),
-                       bp.tools.DotDict(t=bm.arange(args.n_time, dtype=bm.float_),
-                                        i=bm.arange(args.n_time, dtype=bm.int_)),
-                       dyn_vars=net.vars().unique())
+                       DotDict(t=bm.arange(args.n_time, dtype=bm.float_),
+                               i=bm.arange(args.n_time, dtype=bm.int_)),
+                       child_objs=net)
 
   # loss function
   @bm.to_object(child_objs=net)
   def loss_fun(X, Y, fit=True):
-    fr = bm.mean(inference_fun(X, fit), axis=0)
+    fr = bm.max(inference_fun(X, fit), axis=0)
     ys_onehot = bm.one_hot(Y, 10, dtype=bm.float_)
     l = bp.losses.mean_squared_error(fr, ys_onehot)
     n = bm.sum(fr.argmax(1) == Y)
     return l, n
 
-  predict_loss_fun = bm.jit(partial(loss_fun, fit=True), dyn_vars=loss_fun.vars().unique())
+  predict_loss_fun = bm.jit(partial(loss_fun, fit=True), child_objs=loss_fun)
 
   grad_fun = bm.grad(loss_fun, grad_vars=net.train_vars().unique(), has_aux=True, return_value=True)
 
@@ -242,7 +242,7 @@ def train_fun(X, Y):
         'train_acc': train_acc,
         'test_acc': test_acc,
       }
-      bp.checkpoints.save(out_dir, states, epoch_i)
+      # bp.checkpoints.save(out_dir, states, epoch_i)
 
   # inference
   state_dict = bp.checkpoints.load(out_dir)

examples/training_snn_models/mnist_lif_readout.py

@@ -65,9 +65,10 @@ def update(self, p, x):
 
 @bm.to_object(child_objs=(net, encoder))
 def loss_fun(xs, ys):
-  net.reset_state(xs.shape[0])
+  net.reset_state(batch_size=xs.shape[0])
   xs = encoder(xs, num_step=args.T)
-  shared = bm.form_shared_args(num_step=xs.shape[0])
+  # shared arguments for looping over time
+  shared = bm.shared_args_over_time(num_step=args.T)
   outs = bm.for_loop(net, (shared, xs))
   out_fr = bm.mean(outs, axis=0)
   ys_onehot = bm.one_hot(ys, 10, dtype=bm.float_)
@@ -140,7 +141,7 @@ def train(xs, ys):
 state_dict = bp.checkpoints.load(out_dir)
 net.load_state_dict(state_dict['net'])
 
-runner = bp.DSRunner(net, time_major=True)
+runner = bp.DSRunner(net, data_first_axis='T')
 correct_num = 0
 for i in range(0, x_test.shape[0], 512):
   X = encoder(x_test[i: i + 512], num_step=args.T)