Use Flax modules as a part of the BrainPy program#
import brainpy as bp
import brainpy.math as bm
import brainpy_datasets as bd
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
from functools import partial
from flax import linen as nn
bm.set(mode=bm.training_mode, dt=1.)
bp.__version__
'2.8.0'
In this example, we use the Flax, a library used for deep neural networks, to define a convolutional neural network (CNN). The, we integrate this CNN model into our RNN model which defined by BrainPy’s syntax.
Here, we first use flax to define a CNN network.
class CNN(nn.Module):
"""A CNN model implemented by using Flax."""
@nn.compact
def __call__(self, x):
x = nn.Conv(features=32, kernel_size=(3, 3))(x)
x = nn.relu(x)
x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
x = nn.Conv(features=64, kernel_size=(3, 3))(x)
x = nn.relu(x)
x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
x = x.reshape((x.shape[0], -1)) # flatten
x = nn.Dense(features=256)(x)
x = nn.relu(x)
return x
Then, we define an RNN model by using our BrainPy interface. Note here, the Flax module is used as a module at one single step.
class Network(bp.DynamicalSystemNS):
def __init__(self):
super(Network, self).__init__()
self.cnn = bp.dnn.FromFlax(
CNN(), # the model
bm.ones([1, 4, 28, 1]) # an example of the input used to initialize the model parameters
)
self.rnn = bp.dyn.GRUCell(256, 100)
self.linear = bp.dnn.Dense(100, 10)
def update(self, x):
x = self.cnn(x)
x = self.rnn(x)
x = self.linear(x)
return x
We initialize the network, optimizer, loss function, and BP trainer.
net = Network()
opt = bp.optim.Momentum(0.1)
We get the MNIST dataset.
data = bd.vision.MNIST(r'D:\data', download=True)
data.data = data.data.reshape(-1, 7, 4, 28, 1) / 255
def get_data(batch_size):
key = bm.random.split_key()
data.data = bm.random.permutation(data.data, key=key)
data.targets = bm.random.permutation(data.targets, key=key)
for i in range(0, len(data), batch_size):
yield data.data[i: i + batch_size], data.targets[i: i + batch_size]
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Failed to download (trying next):
HTTP Error 404: Not Found
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to D:\data/MNIST/raw/train-images-idx3-ubyte.gz
0%| | 0/9912422 [00:00<?, ?it/s]
0%| | 17408/9912422 [00:00<02:03, 79961.59it/s]
1%| | 69632/9912422 [00:00<00:54, 180925.71it/s]
2%|▏ | 174080/9912422 [00:00<00:30, 321353.50it/s]
4%|▍ | 382976/9912422 [00:00<00:16, 576666.70it/s]
8%|▊ | 783360/9912422 [00:01<00:09, 953026.08it/s]
16%|█▋ | 1612800/9912422 [00:01<00:04, 2066614.52it/s]
26%|██▌ | 2545664/9912422 [00:01<00:02, 3439782.01it/s]
33%|███▎ | 3244032/9912422 [00:01<00:01, 4049469.83it/s]
43%|████▎ | 4298752/9912422 [00:01<00:01, 5521512.92it/s]
60%|█████▉ | 5912576/9912422 [00:01<00:00, 8117931.21it/s]
76%|███████▋ | 7559168/9912422 [00:01<00:00, 10274547.55it/s]
92%|█████████▏| 9152512/9912422 [00:01<00:00, 11799216.24it/s]
9913344it [00:02, 4772805.98it/s]
Extracting D:\data/MNIST/raw/train-images-idx3-ubyte.gz to D:\data/MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Failed to download (trying next):
HTTP Error 404: Not Found
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz to D:\data/MNIST/raw/train-labels-idx1-ubyte.gz
0%| | 0/28881 [00:00<?, ?it/s]
64%|██████▍ | 18432/28881 [00:00<00:00, 84872.02it/s]
29696it [00:00, 135694.43it/s]
Extracting D:\data/MNIST/raw/train-labels-idx1-ubyte.gz to D:\data/MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
Failed to download (trying next):
HTTP Error 404: Not Found
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz to D:\data/MNIST/raw/t10k-images-idx3-ubyte.gz
0%| | 0/1648877 [00:00<?, ?it/s]
1%| | 17408/1648877 [00:00<00:20, 80823.09it/s]
4%|▍ | 63488/1648877 [00:00<00:10, 158284.46it/s]
11%|█ | 177152/1648877 [00:00<00:04, 319322.28it/s]
22%|██▏ | 368640/1648877 [00:00<00:02, 529145.99it/s]
48%|████▊ | 786432/1648877 [00:01<00:00, 1047358.88it/s]
67%|██████▋ | 1112064/1648877 [00:01<00:00, 1472451.28it/s]
97%|█████████▋| 1604608/1648877 [00:01<00:00, 2061041.40it/s]
1649664it [00:01, 1241268.36it/s]
Extracting D:\data/MNIST/raw/t10k-images-idx3-ubyte.gz to D:\data/MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Failed to download (trying next):
HTTP Error 404: Not Found
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz to D:\data/MNIST/raw/t10k-labels-idx1-ubyte.gz
0%| | 0/4542 [00:00<?, ?it/s]
5120it [00:00, 4272749.00it/s]
Extracting D:\data/MNIST/raw/t10k-labels-idx1-ubyte.gz to D:\data/MNIST/raw
def loss_func(predictions, targets):
logits = bm.max(predictions, axis=1)
loss = bp.losses.cross_entropy_loss(logits, targets)
accuracy = bm.mean(bm.argmax(logits, -1) == targets)
return loss, {'accuracy': accuracy}
Finally, train our defined model by using BPTT.fit() function.
trainer = bp.BPTT(net, loss_fun=loss_func, optimizer=opt, loss_has_aux=True)
trainer.fit(partial(get_data, batch_size=256), num_epoch=10)
Train 0 epoch, use 66.7977 s, loss 0.6844439506530762, accuracy 0.775825560092926
Train 1 epoch, use 65.1492 s, loss 0.11312742531299591, accuracy 0.9657745957374573
Train 2 epoch, use 67.3738 s, loss 0.07630933076143265, accuracy 0.9765348434448242
Train 3 epoch, use 66.9896 s, loss 0.05637252330780029, accuracy 0.9821420311927795
Train 4 epoch, use 66.8185 s, loss 0.04613572731614113, accuracy 0.9852337837219238
Train 5 epoch, use 67.1337 s, loss 0.03939102590084076, accuracy 0.9874777793884277
Train 6 epoch, use 66.9745 s, loss 0.03396125137805939, accuracy 0.9892010688781738
Train 7 epoch, use 67.1286 s, loss 0.028259821236133575, accuracy 0.9911735653877258
Train 8 epoch, use 69.9983 s, loss 0.021847538650035858, accuracy 0.9931516051292419
Train 9 epoch, use 66.5559 s, loss 0.01820436865091324, accuracy 0.9948636889457703