I’ve hit a wall with trying to optimize my code…
I’m trying to type-II maximum likelihood training for exact Gaussian process using Flux to get the hang of things. As much as I’ve tried to optimize my code, I’m struggling to get performance that is comparable to PyTorch on my laptop CPU. For example, training in PyTorch takes about ~75 seconds while training with Flux takes about ~400 seconds. After checking my code avoids obvious problems (I hope), I’ve tried the following:
Setting the
JULIA_NUM_THREADSenvironment variable to the number of available cores.
-
Setting the number of BLAS threads to 12
The most expensive operation in the “forward pass” of the GP is sped up by ~ a factor of 2 when increasing the number of BLAS threads but the training time is increased by ~ 20 seconds.
This makes me think that perhaps the problem is coming from the fact that Flux is not taking advantage of my multiple cpus. After reading posts like this one, I wonder if Flux isn’t intended for this use case? (But maybe this is coming from my own ignorance of Julia+Flux).
I should note that on a previous question it was pointed out that log and exp are slower than they should be but this will be fixed in 1.6. I don’t think that the problems with these functions should be causing the slow down I’m seeing.
Any help that someone could provide would be greatly appreciated here! I’d love to recommend Julia across the board to my lab but I want to make sure I can produce code with comparable performance to PyTorch on problems like the ones we work on first.
For anyone interested, I’ve attached my code below. Note run times were calculated by running
> main(400, 5000) > @time main(400, 5000)
from the REPL.
using LinearAlgebra using Flux using Flux.Optimise: update! using Plots """ Computes the distance between each coordinate for each point in x Returns a matrix whose dimension is n+1 where n is the dimension of x """ function difference(x1::AbstractArray{T,2}, x2::AbstractArray{T,2}) where T s1, s2 = size(x1) t1, t2 = size(x2) return reshape(x1, (s1, s2, 1)) .- reshape(x2, (t1, 1, t2)) end """ Squared exponential kernel with length scale ℓ and σf = exp(logσf) """ struct sekernel{T <: AbstractArray} logℓ::T logσf::T end Flux.@functor sekernel sekernel(ℓ::AbstractFloat, σf::AbstractFloat) = sekernel([ℓ], [σf]) function get_natural_params(k::sekernel) return exp(k.logℓ[1]), exp(k.logσf[1]) end function (k::sekernel)(x1::AbstractArray{T,2}, x2::AbstractArray{T,2}) where T ℓ, σf = get_natural_params(k) sqrd = dropdims(sum(difference(x1, x2).^2, dims=1), dims=1) return σf^2 .* exp.(-1 / (2 * ℓ^2) * sqrd) end """ Gaussian process container x : training inputs, f : training targets, k : kernel struct, logσn : i.i.d noise """ struct GP{U,S <: AbstractArray,T <: AbstractArray,I <: AbstractArray} x::S f::T k::U logσn::I end Flux.@functor GP function GP(x::AbstractArray, f::AbstractArray, logσn::AbstractFloat, k) return GP(x, f, k, [logσn]) end function get_natural_params(g::GP) return exp(g.logσn[1]) end Flux.trainable(g::GP) = (g.logσn, g.k) function (g::GP)(xs::AbstractArray) σn = get_natural_params(g) μ, Σ = infer(g.f, g.k(g.x, g.x) + I .* σn^2, g.k(g.x, xs), g.k(xs, xs)) return (μ, Σ) end """ Function for performing inference with a GP """ function infer(f::AbstractArray{T}, # check this (maybe just AbstractArray{T}?) K11::AbstractArray{T,2}, K12::AbstractArray{T,2}, K22::AbstractArray{T,2}) where T C = cholesky(Symmetric(K11)) # adding 1e-12 for numerical stability μ = K12' * (C \ f) Σ = K22 - K12' * (C \ K12) return μ, Σ end """ Computes the log_mll for a Gaussian process """ function log_mll(f::AbstractArray{T}, K11::AbstractArray{T,2}) where T C = cholesky(Symmetric(K11)) return -0.5 * f' * (C \ f) - 0.5 * logdet(C) - length(f) / 2.0 * log(2.0 * π) end function log_mll(g::GP) σn= get_natural_params(g) return log_mll(g.f, g.k(g.x,g.x) + I .* σn^2) end """ type-ii maximum likelihood training of GP """ function train_gp!(model::GP, epochs::Integer, lr=1e-1) opt = Flux.Optimise.ADAM(lr) ps = Flux.params(model) for i in 1:epochs gs = gradient(ps) do return -1.0*log_mll(model) end update!(opt, ps, gs) if i % 1000 == 0 @info "Epoch $i | len $(exp.(model.k.logℓ)) | σf $(exp.(model.k.logσf)) | σn $(exp.(model.logσn)) | lml $(log_mll(model))" end end end function main(num_data, epochs) x = collect(LinRange(0., 10., num_data)) y_true = x + sin.(x) * 5 .- 10.0 y = y_true + randn(length(x)) * 2.0 model = GP(reshape(x, (1, num_data)), y, -4.0, sekernel(0.1, -4.0)) train_gp!(model, epochs) xt = collect(LinRange(0., 10., 500)) μ, Σ = model(reshape(xt,(1,500)) ) σ = 2 * sqrt.(diag(Σ)) p = scatter(x, y, label="Measurements", alpha=0.5, xlabel="x", ylabel="y") p = plot!(x, y_true, linestyle=:dash, label="Generating Func.", linewidth=2) p = plot!(xt, μ, ribbon=σ, label="Exact GP", linewidth=2) end