Static Neural Network Parameters

We can also allocate neural network parameters using StaticArrays. Therefore we simply need to set the keyword static to true in the NeuralNetwork constructor.

Warning

Static neural network parameters are only supported for dense CPU arrays. AbstractNeuralNetworks defines a type CPUStatic, but does not have equivalent GPU objects.

using AbstractNeuralNetworks
import Random
Random.seed!(123)

backend = AbstractNeuralNetworks.CPUStatic()
input_dim = 2
n_hidden_layers = 100
c = Chain(Dense(input_dim, 10, tanh), Tuple(Dense(10, 10, tanh) for _ in 1:n_hidden_layers)..., Dense(10, 1, tanh))
nn = NeuralNetwork(c, backend)
typeof(nn.params.L1.W)
MMatrix{10, 2, Float64, 20} (alias for StaticArraysCore.MArray{Tuple{10, 2}, Float64, 2, 20})

We can compare different evaluation times:

nn_cpu = changebackend(CPU(), nn)
second_dim = 200
x = rand(input_dim, second_dim)
@time nn(x);
  0.002905 seconds (714 allocations: 3.107 MiB)
@time nn_cpu(x);
  0.002351 seconds (714 allocations: 3.107 MiB)

If we also make the input static, we get:

using StaticArrays
x = @SMatrix rand(input_dim, second_dim)
nn(x);
@time nn(x);
  0.002255 seconds (207 allocations: 1.559 MiB)
@time nn_cpu(x);
  0.002344 seconds (715 allocations: 3.111 MiB)