Neural Network in Julia (Multilayer Perceptron)

I wrote a simple multilayer perceptron in Julia, which seems to work fine on different datasets, e.g. the MNIST dataset with a success rate of about 90% after a few seconds of training. But I would like constructive criticism of the program and I'm very new to the Julia programming language, so feedback on the formatting is very welcome.

This is network.jl:

# network: weighted sums, activations, weights, biases, gradient
mutable struct NN
dims :: Vector{Int}
z :: Vector{Vector{Float64}}
a :: Vector{Vector{Float64}}
w :: Vector{Matrix{Float64}}
b :: Vector{Vector{Float64}}
∇a :: Vector{Vector{Float64}}
∇w :: Vector{Matrix{Float64}}
∇b :: Vector{Vector{Float64}}
Σ∇w :: Vector{Matrix{Float64}}
Σ∇b :: Vector{Vector{Float64}}
end

Data = Vector{Tuple{Vector{Float64}, Vector{Float64}}}

# fill vectors and matrices, needs improvement
function init(dims :: Vector{Int}) :: NN
n = NN([], [], [], [], [], [], [], [], [], [])
n.dims = dims

for i in 1:length(dims)
push!(n.a,   Vector{Float64}(undef, dims[i]))
push!(n.∇a,  Vector{Float64}(undef, dims[i]))
end

for i in 2:length(dims)
push!(n.z,   Vector{Float64}(undef, dims[i]))

push!(n.w,   randn(dims[i], dims[i - 1]))
push!(n.∇w,  Matrix{Float64}(undef, dims[i], dims[i - 1]))
push!(n.Σ∇w, Matrix{Float64}(undef, dims[i], dims[i - 1]))

push!(n.b,   randn(dims[i]))
push!(n.∇b,  Vector{Float64}(undef, dims[i]))
push!(n.Σ∇b, Vector{Float64}(undef, dims[i]))
end

return n
end

# leaky ReLU to avoid dead neurons
ReLU(x)  = max(0.01 * x, x)
ReLU′(x) = x >= 0 ? 1 : 0.01

σ(x)  = 1 / (1 + exp(-x))
σ′(x) = σ(x) * (1 - σ(x))

const act  = σ
const act′ = σ′

function forward!(n :: NN)
for i in 1:length(n.dims) - 1
n.z[i] = n.w[i] * n.a[i] + n.b[i]
n.a[i + 1] = act.(n.z[i])
end
end

cost(output, expected) = sum((output - expected) .^ 2)

function backprop!(n :: NN, expected :: Vector{Float64})
len = length(n.dims)

n.∇a[len] = 2 .* (n.a[len] - expected)

for i in len - 1:-1:1
n.∇a[i] .= 0
n.∇b[i] = act′.(n.z[i]) .* n.∇a[i + 1]
for j in 1:dims[i + 1]
if i != 1
n.∇a[i] += n.w[i][j, :] .* n.∇b[i][j]
end
n.∇w[i][j, :] = n.a[i] .* n.∇b[i][j]
end
end
end

# data: [(input, expected)], only one batch!
function train!(n :: NN, data :: Data) :: Float64
for i in 1:length(n.dims) - 1
n.Σ∇w[i] .= 0
n.Σ∇b[i] .= 0
end

Σcost = 0

for d in data
n.a[1] = d[1]
forward!(n)
Σcost += cost(n.a[length(n.dims)], d[2]) / length(data)

backprop!(n, d[2])
n.Σ∇w .+= n.∇w / length(data)
n.Σ∇b .+= n.∇b / length(data)
end

# play around with factor
n.w -= 5 * n.Σ∇w
n.b -= 5 * n.Σ∇b

return Σcost
end


Example approximating sin(x)

I added the following example of approximating sin(x) using the above network to the GitLab repository. You might want to reduce the number of batches, depending on your computing resources.

Here is examples/sin.jl:

using Printf

include("../network.jl")

dims = [1, 10, 10, 10, 1]
len = length(dims)

n = init(dims)

# batched training data: [[(input, expected)]]
batches = [Data(undef, 5) for i in 1:100000]
for batch in batches
for j in eachindex(batch)
x = rand() * pi * 2
batch[j] = ([x], [sin(x) / 2 + 0.5])
end
end

for batch in batches
@printf "Σcost = %.12f\n" train!(n, batch)
end

println("\nTesting with random values:\n---------------------------")
for i in 1:10
n.a[1][1] = rand() * pi * 2
forward!(n)
expected = sin(n.a[1][1]) / 2 + 0.5
@printf "sin(%.6f) = %.6f | NN: %.6f\n" n.a[1][1] expected n.a[len][1]
end


I also wrote some Pluto.jl notebooks with a friend, that use the above neural network and train it on e.g. the MNIST dataset, if you want to have a look.

Overall this looks very good. The main change I would make is putting a @views on backprop! since slicing in Julia makes a copy by default.

function backprop!(n :: NN, expected :: Vector{Float64})
len = length(n.dims)

n.∇a[len] = 2 .* (n.a[len] - expected)

for i in len - 1:-1:1
n.∇a[i] .= 0
n.∇b[i] = act′.(n.z[i]) .* n.∇a[i + 1]
@views for j in 1:dims[i + 1]
if i != 1
n.∇a[i] += n.w[i][j, :] .* n.∇b[i][j]
end
n.∇w[i][j, :] = n.a[i] .* n.∇b[i][j]
end
end
end