Build Artificial Neural Networks

Artificial neural networks in BrainPy are used to build dynamical systems. Here we only talk about how to build a neural network and how to train it.

The brainpy.simulation.layers module provides various classes representing the layers of a neural network. All of them are subclasses of the brainpy.simulation.layers.Module base class.

import brainpy as bp

import brainpy.simulation.layers as nn
import brainpy.math.jax as bm

Creating a layer

A layer can be created as an instance of a brainpy.layers.Module subclass. For example, a dense layer can be created as follows:

l = nn.Dense(num_hidden=100, num_input=128) 

This will create a dense layer with 100 units, connected to another input layer with 128 dimension.

Creating a network

Chaining layer instances together like this will allow you to specify your desired network structure.

This can be done with inheritance from brainpy.layers.Module,

class MLP(nn.Module):
    def __init__(self, n_in, n_l1, n_l2, n_out):
        super(MLP, self).__init__()
        self.l1 = nn.Dense(num_hidden=n_l1, num_input=n_in)
        self.l2 = nn.Dense(num_hidden=n_l2, num_input=n_l1)
        self.l3 = nn.Dense(num_hidden=n_out, num_input=n_l2)
    def update(self, x):
        x = bm.relu(self.l1(x))
        x = bm.relu(self.l2(x))
        x = self.l3(x)
        return x
mlp1 = MLP(10, 50, 100, 2)

Or using brainpy.layers.Sequential,

mlp2 = nn.Sequential(
    l1=nn.Dense(num_hidden=50, num_input=10),
    l2=nn.Dense(num_hidden=100, num_input=50),
    l3=nn.Dense(num_hidden=2, num_input=100),

Naming a layer

For convenience, you can name a layer by specifying the name keyword argument:

l_hidden = nn.Dense(num_hidden=50, num_input=10, name='hidden_layer')

Initializing parameters

Many types of layers, such as brainpy.layers.Dense, have trainable parameters. These are referred to by short names that match the conventions used in modern deep learning literature. For example, a weight matrix will usually be called w, and a bias vector will usually be b.

When creating a layer with trainable parameters, TrainVar will be created for them and initialized automatically. You can optionally specify your own initialization strategy by using keyword arguments that match the parameter variable names. For example:

l = nn.Dense(num_hidden=50, num_input=10, w=bp.initialize.Normal(0.01))

The weight matrix w of this dense layer will be initialized using samples from a normal distribution with standard deviation 0.01 (see brainpy.initialize for more information).

There are several ways to manually initialize parameters:

  • Tensors

If a tensor variable instance is provided, this is used unchanged as the parameter variable. For example:

w = bm.random.normal(0, 0.01, size=(10, 50))
nn.Dense(num_hidden=50, num_input=10, w=w)
<brainpy.simulation.layers.dense.Dense at 0x23cff9bb910>
  • callable

If a callable is provided (e.g. a function or a brainpy.initialize.Initializer instance), the callable will be called with the desired shape to generate suitable initial parameter values. The variable is then initialized with those values. For example:

nn.Dense(num_hidden=50, num_input=10, w=bp.initialize.Normal(0.01))
<brainpy.simulation.layers.dense.Dense at 0x23cff9bf2b0>

Or, using a custom initialization function:

def init_w(shape):
    return bm.random.normal(0, 0.01, shape)

nn.Dense(num_hidden=50, num_input=10, w=init_w)
<brainpy.simulation.layers.dense.Dense at 0x23cff9ac670>

Some types of parameter variables can also be set to None at initialization (e.g. biases). In that case, the parameter variable will be omitted. For example, creating a dense layer without biases is done as follows:

nn.Dense(num_hidden=50, num_input=10, b=None)
<brainpy.simulation.layers.dense.Dense at 0x23cff99fa30>

Setup a training

Here, we show an example to train MLP to classify the MNIST images.

import numpy as np
import tensorflow as tf

# Data
(X_train, Y_train), (X_test, Y_test) = tf.keras.datasets.mnist.load_data()
num_train, num_test = X_train.shape[0], X_test.shape[0]
num_dim =[1:])
X_train = np.asarray(X_train.reshape((num_train, num_dim)) / 255.0, dtype=bm.float_)
X_test = np.asarray(X_test.reshape((num_test, num_dim)) / 255.0, dtype=bm.float_)
Y_train = np.asarray(Y_train.flatten(), dtype=bm.float_)
Y_test = np.asarray(Y_test.flatten(), dtype=bm.float_)
model = MLP(n_in=num_dim, n_l1=256, n_l2=128, n_out=10)
opt = bm.optimizers.Momentum(lr=1e-3, train_vars=model.train_vars())
gv = bm.grad(lambda X, Y: bm.losses.cross_entropy_loss(model(X), Y),
@bm.function(nodes=(model, opt))
def train(x, y):
    grads, loss = gv(x, y)
    return loss
predict = bm.jit(lambda X: bm.softmax(model(X)), dyn_vars=model.vars())
# Training
num_batch = 128
for epoch in range(30):
  # Train
  loss = []
  sel = np.arange(len(X_train))
  for it in range(0, X_train.shape[0], num_batch):
    l = train(X_train[sel[it:it + num_batch]], Y_train[sel[it:it + num_batch]])

  # Eval
  test_predictions = predict(X_test).argmax(1)
  accuracy = np.array(test_predictions).flatten() == Y_test
  print(f'Epoch {epoch + 1:4d}  Train Loss {np.mean(loss):.3f}  Test Accuracy {100 * np.mean(accuracy):.3f}')
Epoch    1  Train Loss 1.212  Test Accuracy 86.410
Epoch    2  Train Loss 0.467  Test Accuracy 89.810
Epoch    3  Train Loss 0.367  Test Accuracy 90.670
Epoch    4  Train Loss 0.325  Test Accuracy 91.470
Epoch    5  Train Loss 0.298  Test Accuracy 92.220
Epoch    6  Train Loss 0.278  Test Accuracy 92.890
Epoch    7  Train Loss 0.261  Test Accuracy 93.140
Epoch    8  Train Loss 0.248  Test Accuracy 93.530
Epoch    9  Train Loss 0.235  Test Accuracy 93.810
Epoch   10  Train Loss 0.224  Test Accuracy 94.010
Epoch   11  Train Loss 0.214  Test Accuracy 94.100
Epoch   12  Train Loss 0.205  Test Accuracy 94.350
Epoch   13  Train Loss 0.196  Test Accuracy 94.540
Epoch   14  Train Loss 0.189  Test Accuracy 94.680
Epoch   15  Train Loss 0.182  Test Accuracy 94.910
Epoch   16  Train Loss 0.175  Test Accuracy 95.070
Epoch   17  Train Loss 0.169  Test Accuracy 95.190
Epoch   18  Train Loss 0.163  Test Accuracy 95.280
Epoch   19  Train Loss 0.157  Test Accuracy 95.410
Epoch   20  Train Loss 0.153  Test Accuracy 95.570
Epoch   21  Train Loss 0.148  Test Accuracy 95.760
Epoch   22  Train Loss 0.143  Test Accuracy 95.760
Epoch   23  Train Loss 0.139  Test Accuracy 95.930
Epoch   24  Train Loss 0.135  Test Accuracy 95.910
Epoch   25  Train Loss 0.131  Test Accuracy 96.150
Epoch   26  Train Loss 0.128  Test Accuracy 96.110
Epoch   27  Train Loss 0.124  Test Accuracy 96.310
Epoch   28  Train Loss 0.121  Test Accuracy 96.330
Epoch   29  Train Loss 0.118  Test Accuracy 96.410
Epoch   30  Train Loss 0.115  Test Accuracy 96.410