Package: nerva-jax
The nerva-jax libraryNOTE:This page is a stub. The library has not yet been released!This repository contains an implementation of multilayer perceptrons in JAX.
It is part of a group of five Python packages that can be installed via pip:nerva-jaxAn implementation inJAX.nerva-numpyAn implementation inNumPy.nerva-sympyAn implementation inSymPy, used for validation and testing.nerva-tensorflowAn implementation inTensorFlow.nerva-torchAn implementation inPyTorch.The packages can be installed standalone, except fornerva-sympywhich requires
installation of the other four. Each package has its own GitHub repository:https://github.com/wiegerw/nerva-jaxhttps://github.com/wiegerw/nerva-numpyhttps://github.com/wiegerw/nerva-sympyhttps://github.com/wiegerw/nerva-tensorflowhttps://github.com/wiegerw/nerva-torchPurposeThe main purpose of these repositories is on the practical implementation of
neural networks. We aim to achieve the following goals:To provide precise mathematical specifications of the execution of multilayer perceptrons.To provide an overview of the equations for common layers, activation functions and loss functions.To provide easily understandable implementations that match the specifications closely.An important advantage of our implementations is that they
are fully transparent: even the implementation of
backpropagation is provided in an accessible manner. Currently, the scope is
limited to multilayer perceptrons. However, the approach can easily be
generalized to more complex neural network architectures.DocumentationTODOInstallationThe code is available as the PyPI packageCONEstrip.
It can be installed usingpip install nerva-jaxLicensingThe code is available under theBoost Software License 1.0.
Alocal copyis included in the repository.Using the libraryMultilayer perceptronsOur multilayer perceptron class has a straightforward implementation:classMultilayerPerceptron(object):deffeedforward(self,X:Matrix)->Matrix:forlayerinself.layers:X=layer.feedforward(X)returnXdefbackpropagate(self,Y:Matrix,DY:Matrix)->None:forlayerinreversed(self.layers):layer.backpropagate(Y,DY)Y,DY=layer.X,layer.DXdefoptimize(self,lr:float):forlayerinself.layers:layer.optimize(lr)A multilayer perceptron can be constructed like this:M=MultilayerPerceptron()input_size=3072output_size=1024act=AllReLUActivation(0.3)layer=ActivationLayer(input_size,output_size,act)layer.set_optimizer('Momentum(0.9)')layer.set_weights('Xavier')M.layers.append(layer)...An example of a layer is the softmax layer. Note that the feedforward and
backpropagation implementations exactly match with the equations given in the documentation.classSoftmaxLayer(LinearLayer):deffeedforward(self,X:Matrix)->Matrix:...Z=W@X+column_repeat(b,N)Y=softmax_colwise(Z)...defbackpropagate(self,Y:Matrix,DY:Matrix)->None:...DZ=hadamard(Y,DY-row_repeat(diag(Y.T@DY).T,K))DW=DZ@X.TDb=rows_sum(DZ)DX=W.T@DZ...The gradients computed by backpropagation are used to update the parameters of
the neural network using anoptimizer. The user has full control over how to
update them. The composite design pattern is used to achieve this:optimizer_W=MomentumOptimizer(layer.W,layer.DW,0.9)optimizer_b=NesterovOptimizer(layer.b,layer.Db,0.9)layer.optimizer=CompositeOptimizer(optimizer_W,optimizer_b)HereWandbare the weights and bias of a linear layer, andDWandDbare the gradients of these parameters with respect to the loss function.
Other parameters can be learned as well, see for example the SReLU layer.TrainingFor training of an MLP we provide an implementation of stochastic gradient descent
in the filetraining.pythat looks like this:defstochastic_gradient_descent(M:MultilayerPerceptron,epochs:int,loss:LossFunction,learning_rate:LearningRateScheduler,train_loader:DataLoader):forepochinrange(epochs):lr=learning_rate(epoch)for(X,T)intrain_loader:Y=M.feedforward(X)DY=loss.gradient(Y,T)/Y.shape[0]# divide by the number of examplesM.backpropagate(Y,DY)M.optimize(lr)Heretrain_loaderis an object similar to aDataLoaderin PyTorch that splits
a dataset into batches of inputsXand expected outputsT.
Training of an MLP consists of three steps:feedforwardGiven an input batchXand expected outputsT, compute the outputY.backpropagationGiven outputsYand expected outputsT, compute the
gradient of the MLP parameters with respect to a given loss functionL.optimizationUpdate the MLP parameters using the gradient computed in step 2.These steps are performed for each input batch of a dataset, and this process is
repeatedepochtimes.Command line scriptFor convenience, a command line scripttools/mlp.pyis included that can be
used to do a training experiment. An example invocation of this script is
provided inexamples/cifar10.sh:mlp.py--layers="ReLU;ReLU;Linear"\--sizes="3072,1024,512,10"\--optimizers="Momentum(0.9);Momentum(0.9);Momentum(0.9)"\--init-weights="Xavier,Xavier,Xavier"\--batch-size=100\--epochs=10\--loss=SoftmaxCrossEntropy\--learning-rate="Constant(0.01)"\--dataset="cifar10.npz"The output of this script may look like this:$ ./cifar10.sh 
Loading dataset from file ../data/cifar10.npz
epoch   0  lr: 0.01000000  loss: 2.49815798  train accuracy: 0.10048000  test accuracy: 0.10110000  time: 0.00000000s
epoch   1  lr: 0.01000000  loss: 1.64330590  train accuracy: 0.41250000  test accuracy: 0.40890000  time: 11.27521224s
epoch   2  lr: 0.01000000  loss: 1.54620886  train accuracy: 0.44674000  test accuracy: 0.43910000  time: 11.33507117s
epoch   3  lr: 0.01000000  loss: 1.46849191  train accuracy: 0.47462000  test accuracy: 0.46280000  time: 11.30941587s
epoch   4  lr: 0.01000000  loss: 1.40283990  train accuracy: 0.49964000  test accuracy: 0.48370000  time: 11.30728618s
epoch   5  lr: 0.01000000  loss: 1.36808932  train accuracy: 0.51214000  test accuracy: 0.49030000  time: 11.32811917s
epoch   6  lr: 0.01000000  loss: 1.33309329  train accuracy: 0.52786000  test accuracy: 0.49490000  time: 11.27888232s
epoch   7  lr: 0.01000000  loss: 1.31322646  train accuracy: 0.53440000  test accuracy: 0.49800000  time: 11.29106200s
epoch   8  lr: 0.01000000  loss: 1.29327416  train accuracy: 0.53940000  test accuracy: 0.49850000  time: 11.42966828s
epoch   9  lr: 0.01000000  loss: 1.27069771  train accuracy: 0.55052000  test accuracy: 0.50120000  time: 11.30879241s
epoch  10  lr: 0.01000000  loss: 1.24060512  train accuracy: 0.56160000  test accuracy: 0.50690000  time: 11.42121017s
Total training time for the 10 epochs: 113.28471981sValidationWe do not rely on auto differentiation for the backpropagation. Instead, we provide
explicit implementations for it. Since the derivation of these equations is
error-prone, the results have to be validated. A common way to do this is to
rely on gradient checking using numerical approximations. Instead, we are using
symbolic differentation of SymPy. This is provided in thenerva-sympypackage.
Gradient checking of the softmax layer looks like this:# backpropagationDZ=hadamard(Y,DY-row_repeat(diag(Y.T*DY).T,K))DW=DZ*X.TDb=rows_sum(DZ)DX=W.T*DZ# check gradients using symbolic differentiationDW1=diff(loss(Y),w)Db1=diff(loss(Y),b)DX1=diff(loss(Y),x)DZ1=diff(loss(Y),z)self.assertTrue(equal_matrices(DW,DW1))self.assertTrue(equal_matrices(Db,Db1))self.assertTrue(equal_matrices(DX,DX1))self.assertTrue(equal_matrices(DZ,DZ1))An advantage of this approach is that errors can be found in an early stage.
Moreover, the source of such errors can be detected by checking intermediate
results.