Static Neural Network Parameters

We can also allocate neural network parameters using StaticArrays. Therefore we simply need to pass the CPUStatic() backend to the NeuralNetwork constructor.

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(params(nn).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.002184 seconds (714 allocations: 3.109 MiB)

@time nn_cpu(x);

  0.001444 seconds (714 allocations: 3.109 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.001334 seconds (308 allocations: 3.104 MiB)

@time nn_cpu(x);

  0.001517 seconds (714 allocations: 3.109 MiB)