DNN From Scratch

]up; add ImageShow Flux Colors Zygote#master https://github.com/FluxML/ZygoteRules.jl.git ImageMagick FileIO

]precompile

using Colors, ImageShow
import Zygote, Flux

glorot_uniform(dims...) = (rand(Float32, dims...) .- 0.5f0) .* sqrt(24.0f0/sum(dims))

struct Dense{M <: AbstractMatrix, V <: AbstractVector, F <: Function}
  W::M
  b::V
  func::F
end

function Dense(in, out, func = identity)
  Dense(glorot_uniform(out, in), zeros(Float32, out), func)
end

function (a::Dense)(x::AbstractArray)
  a.func.(a.W * x .+ a.b)
end

softmax(xs) = exp.(xs) ./ sum(exp.(xs))

relu(x::Real) = max(zero(x), x)

function crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat; weight = 1)
  -sum(y .* log.(ŷ) .* weight) * 1 // size(y, 2)
end

function forward(network, input)
  result = input
  for layer in network
    result = layer(result)
  end
  return result
end
loss(network, x, y) = crossentropy(forward(network, x), y)
function loss_gradient(network, input, target)
  # first index, to get gradient of first argument
  Zygote.gradient(loss, network, input, target)[1]
end

function apply_gradient!(a, b::Nothing, optimizer = nothing)
  return nothing
end

function apply_gradient!(a, b::NamedTuple, optimizer = nothing)
  for field in propertynames(b)
    apply_gradient!(getfield(a, field), getfield(b, field))
  end
  return nothing
end
function apply_gradient!(a::Tuple, b, optimizer = nothing)
  for (alayer, blayer) in zip(a, b)
    apply_gradient!(alayer, blayer)
  end
  return nothing
end

function apply_gradient!(a::AbstractArray, b::AbstractArray, optimizer = nothing)
  a .-= 0.1 .* b
  return nothing
end

"""
    ADAM(η = 0.001, β = (0.9, 0.999))

[ADAM](https://arxiv.org/abs/1412.6980v8) optimiser.
"""
mutable struct ADAM
  eta::Float64
  beta::Tuple{Float64,Float64}
  state::IdDict
end

ADAM(η = 0.001, β = (0.9, 0.999)) = ADAM(η, β, IdDict())

function apply_gradient!(x::AbstractArray, Δ::AbstractArray, o::ADAM)
  η, β = o.eta, o.beta
  mt, vt, βp = get!(o.state, x, (zero(x), zero(x), β))
  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. vt = β[2] * vt + (1 - β[2]) * Δ^2
  @. Δ =  mt / (1 - βp[1]) / (√(vt / (1 - βp[2])) + ϵ) * η
  o.state[x] = (mt, vt, βp .* β)
  return Δ
end

function test(n)
  img = X[1:28^2, n:n]
  predict = Tuple(argmax(forward(network, img)))[1] - 1
  @show predict
  save("/results/test.png", Gray.(reshape(img, (28, 28))))
  return nothing
end

function train!(network, X, Y, optimizer = nothing, epochs = 100)
  for epoch in 1:epochs
    grad = loss_gradient(network, X, Y)
    apply_gradient!(network, grad, optimizer)
    @show epoch
  end
end

network = (
  Dense(28^2, 32, relu),
  Dense(32, 10),
  softmax
)
imgs = Flux.Data.MNIST.images()
labels = Flux.Data.MNIST.labels()
Y = Flux.onehotbatch(labels, 0:9)
X = Float32.(hcat(float.(reshape.(imgs, :))...))
train!(network, X, Y, ADAM())

test(23)