Closed Jarrod-Angove closed 2 months ago
Hmmm I'm really not sure what Zygote is doing here -- it has an annoying habit of breaking when you don't expect it to...
I would suggest trying with Tapir.jl. Could you please see if the following works for you?
using Stheno, Tapir, Random, Optim, ParameterHandling
l1 = 0.4
s1 = 0.2
l2 = 5.0
s2 = 1.0
g = @gppp let
f1 = s1 * stretch(GP(Matern52Kernel()), 1 / l1)
f2 = s2 * stretch(GP(SEKernel()), 1 / l2)
f3 = f1 + f2
end;
x = GPPPInput(:f3, collect(range(-5.0, 5.0; length=100)));
σ²_n = 0.02;
fx = g(x, σ²_n);
y = rand(fx);
θ = (
# Short length-scale and small variance.
l1 = positive(0.4),
s1 = positive(0.2),
# Long length-scale and larger variance.
l2 = positive(5.0),
s2 = positive(1.0),
# Observation noise variance -- we'll be learning this as well. Constrained to be
# at least 1e-3.
s_noise = positive(0.1, exp, 1e-3),
)
θ_flat_init, unflatten = flatten(θ);
function build_model(θ::NamedTuple)
return @gppp let
f1 = θ.s1 * stretch(GP(SEKernel()), 1 / θ.l1)
f2 = θ.s2 * stretch(GP(SEKernel()), 1 / θ.l2)
f3 = f1 + f2
end
end
function nlml(θ::NamedTuple)
f = build_model(θ)
return -logpdf(f(x, θ.s_noise + 1e-6), y)
end
# Define objective function, check it runs, and compute a gradient to check that works.
obj(x) = nlml(ParameterHandling.value(unflatten(x)))
obj(θ_flat_init)
rule = Tapir.build_rrule(obj, θ_flat_init)
Tapir.value_and_gradient!!(rule, obj, θ_flat_init)
# Run optimisation.
results = Optim.optimize(
obj,
θ->Tapir.value_and_gradient!!(rule, obj, θ)[2][2],
θ_flat_init + randn(length(θ_flat_init)),
BFGS(),
Optim.Options(show_trace=true,);
inplace=false,
)
Thanks for the quick response @willtebbutt !
I've run the suggested code but the Tapir.build_rrule
function is failing:
julia> rule = Tapir.build_rrule(obj, θ_flat_init)
ERROR: MethodError: no method matching tangent_field_type(::Type{ParameterHandling.var"#unflatten_to_NamedTuple#15"{…}}, ::Int64)
The applicable method may be too new: running in world age 31913, while current world is 34727.
Closest candidates are:
tangent_field_type(::Type{P}, ::Int64) where P (method too new to be called from this world context.)
@ Tapir ~/.julia/packages/Tapir/7eB9t/src/tangents.jl:282
Edit: I thought it may be an issue with the ParameterHandling
package, so I changed theta to a vector and tried Zygote
, Tapir
, and ForwardDiff
again;
θ = [
0.4,
0.2,
5.0,
1.0,
0.1,
]
Zygote still fails with the same type error. Tapir fails similarly, but this time the build_rrule
function works and the value_and_gradient!!
function fails:
ERROR: MethodError: no method matching tangent_field_type(::Type{Stheno.GaussianProcessProbabilisticProgramme{@NamedTuple{…}}}, ::Int64)
The applicable method may be too new: running in world age 31913, while current world is 34966.
Interestingly, by not using ParameterHandling
, the ForwardDiff
package is able to compute the gradient:
julia> ForwardDiff.gradient(obj, θ)
5-element Vector{Float64}:
-87.61487924857141
-182.63178639865728
5.147333884587856
17.701534645018494
226.10513088590136
Hmm could you show me the output of Pkg.status()
. It might be that you need a clean install. Also, what version of Julia are you on?
Yep... It looks like there is something wrong with my project. I ran the code again in a clean install and it appears to work fine. I have no idea what could be causing this, but I'll try rebuilding everything piece by piece until I find the conflict.
For the sake of completeness:
My Julia version is 1.10.4. and here is my Pkg.status()
:
Project DilPredict v0.1.0
Status `~/Documents/grad_school/thesis/DilPredict/Project.toml`
⌃ [99985d1d] AbstractGPs v0.5.9
[8bb1440f] DelimitedFiles v1.9.1
[39dd38d3] Dierckx v0.5.3
[f6369f11] ForwardDiff v0.10.36
[033835bb] JLD2 v0.4.48
[ec8451be] KernelFunctions v0.10.63
[429524aa] Optim v1.9.4
[2412ca09] ParameterHandling v0.5.0
[91a5bcdd] Plots v1.40.4
[c46f51b8] ProfileView v1.7.2
[8188c328] Stheno v0.8.2
[07d77754] Tapir v0.2.20
[e88e6eb3] Zygote v0.6.70
[37e2e46d] LinearAlgebra
[de0858da] Printf
[9a3f8284] Random
[10745b16] Statistics v1.10.0
[fa267f1f] TOML v1.0.3
Sorry for the trouble!
Edit: For some reason Pkg had been fetching a super outdated version of AbstractGPs... All I needed to do was update :')
No trouble at all -- happy to have been able to help!
I'm quite new to this package, so please forgive me if this is a mistake on my end, but I've been having some difficulties getting Optim AD methods working with Stheno. Even after directly copying the BFGS example from the docs, I end up with the following error:
This occurs if I copy and past the tutorial in the docs directly into the repl, or if I paste everything as a single function:
This issue does not occur if I do not provide a gradient to optim, as I believe it defaults to a finite different method to calculate the gradient. Similarly, using
NelderMead()
allows it to run without issue.Here is the stack trace: