upgrade examples

Author: chaoming0625

examples/ANN_models/mnist-cnn.py

@@ -2,6 +2,7 @@
 
 import brainpy as bp
 import brainpy.math as bm
+import brainpy_datasets as bd
 
 
 class FeedForwardModel(bp.DynamicalSystem):
@@ -25,7 +26,7 @@ def update(self, s, x):
 
 
 # train dataset
-train_dataset = bp.datasets.MNIST(root='./data', train=True, download=True)
+train_dataset = bd.vision.MNIST(root='./data', split='train', download=True)
 x_train = bm.array(train_dataset.data, dtype=bm.dftype())
 x_train = x_train.reshape(x_train.shape + (1,)) / 255
 y_train = bm.array(train_dataset.targets, dtype=bm.ditype())
@@ -41,7 +42,7 @@ def update(self, s, x):
 trainer.fit([x_train, y_train], num_epoch=2, batch_size=64)
 
 # test dataset
-test_dataset = bp.datasets.MNIST(root='./data', train=False, download=True)
+test_dataset = bd.vision.MNIST(root='./data', split='train', download=True)
 x_test = bm.array(test_dataset.data, dtype=bm.dftype())
 x_test = x_test.reshape(x_test.shape + (1,)) / 255
 y_test = bm.array(test_dataset.targets, dtype=bm.ditype())

examples/training/Song_2016_EI_RNN.py

@@ -154,7 +154,7 @@ def cell(self, x, h):
   def readout(self, h):
     return h @ self.w_ro + self.b_ro
 
-  def make_update(self, h: bm.JaxArray, o: bm.JaxArray):
+  def make_update(self, h: bm.Array, o: bm.Array):
     def f(x):
       h.value = self.cell(x, h.value)
       o.value = self.readout(h.value[:, :self.e_size])
@@ -164,7 +164,7 @@ def f(x):
 
   def predict(self, xs):
     self.h[:] = 0.
-    return bm.for_loop(self.make_update(self.h, self.o), self.vars(), xs)
+    return bm.for_loop(self.make_update(self.h, self.o), xs, dyn_vars=self.vars())
 
   def loss(self, xs, ys):
     hs, os = self.predict(xs)
@@ -191,19 +191,19 @@ def loss(self, xs, ys):
 
 # %%
 # gradient function
-grad_f = bm.grad(net.loss,
-                 dyn_vars=net.vars(),
-                 grad_vars=net.train_vars().unique(),
-                 return_value=True)
+grad = bm.grad(net.loss,
+               dyn_vars=net.vars().unique(),
+               grad_vars=net.train_vars().unique(),
+               return_value=True)
 
 
 # %%
 @bm.jit
-@bm.function(nodes=(net, opt))
+@bm.to_object(child_objs=(grad, opt))  # add nodes and vars used
 def train(xs, ys):
-  grads, loss = grad_f(xs, ys)
+  grads, l = grad(xs, ys)
   opt.update(grads)
-  return loss
+  return l
 
 
 # %% [markdown]

examples/training/SurrogateGrad_lif-ANN-style.py

@@ -10,8 +10,8 @@
 
 import time
 
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
 from matplotlib.gridspec import GridSpec
 
 import brainpy as bp
@@ -94,41 +94,34 @@ def print_classification_accuracy(output, target):
 print_classification_accuracy(out, y_data)
 
 
+@bm.to_object(child_objs=net, dyn_vars=rng)  # add nodes and vars used in this function
 def loss():
   key = rng.split_key()
-  X = rng.permutation(x_data, key=key)
-  Y = rng.permutation(y_data, key=key)
+  X = bm.random.permutation(x_data, key=key)
+  Y = bm.random.permutation(y_data, key=key)
   looper = bp.dyn.DSRunner(net, numpy_mon_after_run=False, progress_bar=False)
   predictions = looper.run(inputs=X, inputs_are_batching=True, reset_state=True)
   predictions = bm.max(predictions, axis=1)
   return bp.losses.cross_entropy_loss(predictions, Y)
 
 
-f_grad = bm.grad(loss,
-                 grad_vars=net.train_vars().unique(),
-                 dyn_vars=net.vars().unique() + {'rng': rng},
-                 return_value=True)
-f_opt = bp.optim.Adam(lr=2e-3, train_vars=net.train_vars().unique())
+grad = bm.grad(loss, grad_vars=loss.train_vars().unique(), return_value=True)
+optimizer = bp.optim.Adam(lr=2e-3, train_vars=net.train_vars().unique())
 
 
+@bm.to_object(child_objs=(grad, optimizer))  # add nodes and vars used in this function
 def train(_):
-  grads, l = f_grad()
-  f_opt.update(grads)
+  grads, l = grad()
+  optimizer.update(grads)
   return l
 
 
-f_train = bm.make_loop(
-  train,
-  dyn_vars=f_opt.vars() + net.vars() + {'rng': rng},
-  has_return=True
-)
-
 # train the network
 net.reset_state(num_sample)
 train_losses = []
 for i in range(0, 3000, 100):
   t0 = time.time()
-  _, ls = f_train(bm.arange(i, i + 100, 1))
+  ls = bm.for_loop(train, operands=bm.arange(i, i + 100, 1))
   print(f'Train {i + 100} epoch, loss = {bm.mean(ls):.4f}, used time {time.time() - t0:.4f} s')
   train_losses.append(ls)
 

examples/training/SurrogateGrad_lif.py

@@ -90,12 +90,12 @@ def print_classification_accuracy(output, target):
 # Before training
 runner = bp.dyn.DSRunner(net, monitors={'r.spike': net.r.spike, 'r.membrane': net.r.V})
 out = runner.run(inputs=x_data, inputs_are_batching=True, reset_state=True)
-# plot_voltage_traces(runner.mon.get('r.membrane'), runner.mon.get('r.spike'))
-# plot_voltage_traces(out)
+plot_voltage_traces(runner.mon.get('r.membrane'), runner.mon.get('r.spike'))
+plot_voltage_traces(out)
 print_classification_accuracy(out, y_data)
 
 
-@bm.function(nodes=net, dyn_vars=rng)  # add nodes and vars used in this function
+@bm.to_object(child_objs=net, dyn_vars=rng)  # add nodes and vars used here
 def loss():
   key = rng.split_key()
   X = bm.random.permutation(x_data, key=key)
@@ -106,14 +106,11 @@ def loss():
   return bp.losses.cross_entropy_loss(predictions, Y)
 
 
-grad = bm.grad(loss,
-               grad_vars=loss.train_vars().unique(),
-               dyn_vars=loss.vars().unique(),
-               return_value=True)
+grad = bm.grad(loss, grad_vars=loss.train_vars().unique(), return_value=True)
 optimizer = bp.optim.Adam(lr=2e-3, train_vars=net.train_vars().unique())
 
 
-@bm.function(nodes=(grad, optimizer))  # add nodes and vars used in this function
+@bm.to_object(child_objs=(grad, optimizer))  # add nodes and vars used here
 def train(_):
   grads, l = grad()
   optimizer.update(grads)
@@ -125,7 +122,7 @@ def train(_):
 train_losses = []
 for i in range(0, 3000, 100):
   t0 = time.time()
-  ls = bm.for_loop(train, dyn_vars=train.vars().unique(), operands=bm.arange(i, i + 100, 1))
+  ls = bm.for_loop(train, operands=bm.arange(i, i + 100, 1))
   print(f'Train {i + 100} epoch, loss = {bm.mean(ls):.4f}, used time {time.time() - t0:.4f} s')
   train_losses.append(ls)
 

examples/training/SurrogateGrad_lif_fashion_mnist.py

@@ -14,6 +14,7 @@
 
 import brainpy as bp
 import brainpy.math as bm
+import brainpy_datasets as bd
 
 
 class SNN(bp.dyn.Network):
@@ -167,7 +168,7 @@ def loss_fun(predicts, targets):
   )
   trainer.fit(lambda: sparse_data_generator(x_data, y_data, batch_size, nb_steps, nb_inputs),
               num_epoch=nb_epochs)
-  return trainer.train_losses
+  return trainer.get_hist_metric('fit')
 
 
 def compute_classification_accuracy(model, x_data, y_data, batch_size=128, nb_steps=100, nb_inputs=28 * 28):
@@ -198,16 +199,8 @@ def get_mini_batch_results(model, x_data, y_data, batch_size=128, nb_steps=100,
 
 # load the dataset
 root = r"D:\data\fashion-mnist"
-train_dataset = bp.datasets.FashionMNIST(root,
-                                         train=True,
-                                         transform=None,
-                                         target_transform=None,
-                                         download=True)
-test_dataset = bp.datasets.FashionMNIST(root,
-                                        train=False,
-                                        transform=None,
-                                        target_transform=None,
-                                        download=True)
+train_dataset = bd.vision.FashionMNIST(root, split='train', download=True)
+test_dataset = bd.vision.FashionMNIST(root, split='test', download=True)
 
 # Standardize data
 x_train = np.array(train_dataset.data, dtype=bm.dftype())
@@ -237,7 +230,7 @@ def get_mini_batch_results(model, x_data, y_data, batch_size=128, nb_steps=100,
 
 nb_plt = 4
 gs = GridSpec(1, nb_plt)
-fig = plt.figure(figsize=(7, 3), dpi=150)
+plt.figure(figsize=(7, 3), dpi=150)
 for i in range(nb_plt):
   plt.subplot(gs[i])
   plt.imshow(bm.as_numpy(spikes[i]).T, cmap=plt.cm.gray_r, origin="lower")

examples/training/Sussillo_Abbott_2009_FORCE_Learning.py

@@ -5,7 +5,9 @@
 # %% [markdown]
 # Implementation of the paper:
 #
-# - Sussillo, David, and Larry F. Abbott. "Generating coherent patterns of activity from chaotic neural networks." Neuron 63, no. 4 (2009): 544-557.
+# - Sussillo, David, and Larry F. Abbott. "Generating coherent patterns
+#   of activity from chaotic neural networks."
+#   Neuron 63, no. 4 (2009): 544-557.
 
 # %%
 import brainpy as bp
@@ -46,7 +48,7 @@ def __init__(self, num_input, num_hidden, num_output,
     self.w_rr = g * bm.random.normal(size=(num_hidden, num_hidden)) / bm.sqrt(num_hidden)
     self.w_or = bm.random.normal(size=(num_output, num_hidden))
     w_ro = bm.random.normal(size=(num_hidden, num_output)) / bm.sqrt(num_hidden)
-    self.w_ro = bm.Variable(w_ro)
+    self.w_ro = bm.Variable(w_ro)  # dynamically change this weight
 
     # variables
     self.h = bm.Variable(bm.random.normal(size=num_hidden) * 0.5)  # hidden
@@ -62,6 +64,7 @@ def update(self, x):
     self.h += self.dt / self.tau * dhdt
     self.r.value = bm.tanh(self.h)
     self.o.value = bm.dot(self.r, self.w_ro)
+    return self.r.value, self.o.value
 
   def rls(self, target):
     # update the inverse correlation matrix
@@ -75,17 +78,14 @@ def rls(self, target):
     self.w_ro += dw
 
   def simulate(self, xs):
-    f = bm.make_loop(self.update, dyn_vars=self.vars(), out_vars=[self.r, self.o])
-    return f(xs)
+    return bm.for_loop(self.update, dyn_vars=self.vars(), operands=xs)
 
   def train(self, xs, targets):
-    def _f(x):
-      input, target = x
-      self.update(input)
+    def _f(x, target):
+      r, o = self.update(x)
       self.rls(target)
-
-    f = bm.make_loop(_f, dyn_vars=self.vars(), out_vars=[self.r, self.o])
-    return f([xs, targets])
+      return r, o
+    return bm.for_loop(_f, dyn_vars=self.vars(), operands=[xs, targets])
 
 
 # %%

examples/training/integrator_rnn.py

@@ -12,9 +12,7 @@
 num_batch = 128
 
 
-@partial(bm.jit,
-         dyn_vars=bp.TensorCollector({'a': bm.random.DEFAULT}),
-         static_argnames=['batch_size'])
+@partial(bm.jit, dyn_vars=bm.random.DEFAULT, static_argnames=['batch_size'])
 def build_inputs_and_targets(mean=0.025, scale=0.01, batch_size=10):
   # Create the white noise input
   sample = bm.random.normal(size=(batch_size, 1, 1))
@@ -31,14 +29,15 @@ def train_data():
     yield build_inputs_and_targets(batch_size=num_batch)
 
 
-class RNN(bp.dyn.DynamicalSystem):
+class RNN(bp.DynamicalSystem):
   def __init__(self, num_in, num_hidden):
     super(RNN, self).__init__()
-    self.rnn = bp.layers.VanillaRNN(num_in, num_hidden, train_state=True)
+    self.rnn = bp.layers.RNNCell(num_in, num_hidden, train_state=True)
     self.out = bp.layers.Dense(num_hidden, 1)
 
   def update(self, sha, x):
-    return self.out(sha, self.rnn(sha, x))
+    return self.out(sha,
+                    self.rnn(sha, x))
 
 
 model = RNN(1, 100)
@@ -58,11 +57,10 @@ def loss(predictions, targets, l2_reg=2e-4):
 # create a trainer
 trainer = bp.train.BPTT(model, loss_fun=loss, optimizer=opt)
 trainer.fit(train_data,
-            batch_size=num_batch,
             num_epoch=30,
             num_report=200)
 
-plt.plot(bm.as_numpy(trainer.train_losses))
+plt.plot(bm.as_numpy(trainer.get_hist_metric()))
 plt.show()
 
 model.reset_state(1)