Custom Loss Function in Flux.jl - julia

Question

I am trying to implement a model with a custom loss function in the Flux.jl package. I include the code for a simplified model, but the error stays the same.
I have an interpolator which takes a scalar value and returns a 2x2 matrix. The goal of my model is to use 3 observations to find the best point to evaluate the interpolator at. For this I wrote a custom loss function that computes the suggested evalutation_point and evaluates the interpolator at this point. Then the interpolated result is compared to the true solution from the dataset.
using Flux, Zygote
using LinearAlgebra
using Interpolations
##
# create interpolator
x = LinRange(0,1,10)
y = [rand(2,2) for i in 1:10]
itp = interpolate(y, BSpline(Linear())) |> i -> scale(i, x)
# create training set
training_set = [(rand(3), rand(2,2)) for i in 0:0.2:1]
#build the model
model = Chain(Dense(3,1),i-> clamp(i[1],0,1))
opt = Descent()
ps = Flux.params(model)
function loss(evaluation_point, solution)
interpolated = itp(model(evaluation_point))
return norm(interpolated - solution)
end
# training NOK
n_epochs = 100
for epoch in 1:n_epochs
Flux.train!(loss, ps, training_set, opt)
println(sum([loss_fnc(i[1],i[2]) for i in training_set]))
end
This returns the following error:
ERROR: DimensionMismatch("matrix A has dimensions (2,2), vector B has length 1")
Stacktrace:
[1] generic_matvecmul!(C::Vector{Matrix{Float64}}, tA::Char, A::Matrix{Float64}, B::StaticArrays.SVector{1, Matrix{Float64}}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
# LinearAlgebra C:\Users\thega\AppData\Local\Programs\Julia-1.7.2\share\julia\stdlib\v1.7\LinearAlgebra\src\matmul.jl:713
[2] mul!
# C:\Users\thega\AppData\Local\Programs\Julia-1.7.2\share\julia\stdlib\v1.7\LinearAlgebra\src\matmul.jl:81 [inlined]
[3] mul!
# C:\Users\thega\AppData\Local\Programs\Julia-1.7.2\share\julia\stdlib\v1.7\LinearAlgebra\src\matmul.jl:275 [inlined]
[4] *
# C:\Users\thega\AppData\Local\Programs\Julia-1.7.2\share\julia\stdlib\v1.7\LinearAlgebra\src\matmul.jl:51 [inlined]
[5] interpolate_pullback
# C:\Users\thega\.julia\packages\Interpolations\Glp9h\src\chainrules\chainrules.jl:13 [inlined]
[6] ZBack
# C:\Users\thega\.julia\packages\Zygote\H6vD3\src\compiler\chainrules.jl:204 [inlined]
[7] Pullback
# c:\Users\thega\Desktop\Question\main.jl:21 [inlined]
[8] (::typeof(∂(loss)))(Δ::Float64)
# Zygote C:\Users\thega\.julia\packages\Zygote\H6vD3\src\compiler\interface2.jl:0
[9] #212
# C:\Users\thega\.julia\packages\Zygote\H6vD3\src\lib\lib.jl:203 [inlined]
[10] #1750#back
# C:\Users\thega\.julia\packages\ZygoteRules\AIbCs\src\adjoint.jl:67 [inlined]
[11] Pullback
# C:\Users\thega\.julia\packages\Flux\0c9kI\src\optimise\train.jl:102 [inlined]
[12] (::typeof(∂(λ)))(Δ::Float64)
# Zygote C:\Users\thega\.julia\packages\Zygote\H6vD3\src\compiler\interface2.jl:0
[13] (::Zygote.var"#93#94"{Params, typeof(∂(λ)), Zygote.Context})(Δ::Float64)
# Zygote C:\Users\thega\.julia\packages\Zygote\H6vD3\src\compiler\interface.jl:357
[14] gradient(f::Function, args::Params)
# Zygote C:\Users\thega\.julia\packages\Zygote\H6vD3\src\compiler\interface.jl:76
[15] macro expansion
# C:\Users\thega\.julia\packages\Flux\0c9kI\src\optimise\train.jl:101 [inlined]
[16] macro expansion
# C:\Users\thega\.julia\packages\Juno\n6wyj\src\progress.jl:134 [inlined]
[17] train!(loss::Function, ps::Params, data::Vector{Tuple{Vector{Float64}, Matrix{Float64}}}, opt::Descent; cb::Flux.Optimise.var"#40#46")
# Flux.Optimise C:\Users\thega\.julia\packages\Flux\0c9kI\src\optimise\train.jl:99
[18] train!(loss::Function, ps::Params, data::Vector{Tuple{Vector{Float64}, Matrix{Float64}}}, opt::Descent)
# Flux.Optimise C:\Users\thega\.julia\packages\Flux\0c9kI\src\optimise\train.jl:97
[19] top-level scope
# c:\Users\thega\Desktop\Question\main.jl:28
So something about a dimension mismatch, but the evaluation of the loss function works fine.
loss(training_set[1][1], training_set[1][2])
I play around a bit and found that the problem is the gradient computation:
gradient(loss , training_set[1][1], training_set[1][2])

Answer 1

The problem can be found in the training set, with your provided example, check this out:
julia> training_set[1][1]
3-element Vector{Float64}:
0.5093876656425886
0.05272770864628318
0.7651982428671759
julia> training_set[1][2]
2×2 Matrix{Float64}:
0.0691616 0.55414
0.5153 0.654379
For the model, the input is: as x a 2-element vector, and the model should learn to return a 2x2 matrix as y. However, your model does not do that:
julia> model(training_set[1][1])
0.6585413f0 (tracked)
It only returns an instance, due to the definition of the model, which in this case model = Chain(Dense(3,1),i-> clamp(i[1],0,1)) turns to be only a Chain(Dense(3, 1), #7), which means that has a 3 element vector as an input, and returns 1 (and only 1) instance.
Solutions:
redefine your y as a 1 element output for each x
redefine your model (it will be more complicated as you want to return a 2x2 matrix). An example of this would be the following model:
julia> model = Chain(
Dense(3, 4),
x -> reshape(x, (2, 2))
)
But then, you should figure out how to adapt your interpolation code to work

Answer 2

I was not able to fix the problem. My guess is that Interpolations.jl is not compatible with Zygote.jl. A possible workaround I found, was writing a custom interpolations class and function. I include a working example if anyone is interested:
using Flux, Zygote
using LinearAlgebra
using Interpolations
# create a custom linear splines class
struct CustomInterpolator
x::Vector
y::Vector
function CustomInterpolator(x,y)
#assert issorted(x)
return new(x,y)
end
end
function custom_interpolate(citp::CustomInterpolator, x::Number)
left_value, right_value = 0, 0
left_index, right_index = 1, 1
# check bound
if x > citp.x[end] || x < citp.x[1]
#error "Out of bounds"
throw(DomainError(x))
end
#find the right indices
for (i,v) in enumerate(citp.x)
if left_value > x
right_value = v
right_index = i
break
end
left_value = v
left_index = i
end
# do a linear inter interpolation between the two selected indices
interpolated_value = (1 - (x - left_value)/(right_value - left_value)) * citp.y[left_index] + (x - left_value)/(right_value - left_value) * citp.y[right_index]
return interpolated_value
end
##
# create custom interpolator
x = LinRange(0,1,2)
y = [zeros(2,2), ones(2,2)]
citp = CustomInterpolator(x,y)
# create training set
training_set = [(ones(3)*i, ones(2,2) - i*ones(2,2)) for i in 0:0.2:1]
#build the model
model = Chain(Dense(3,3), Dense(3,1), i-> clamp(i[1],0,1), i->custom_interpolate(citp,i))
opt = ADAM()
ps = Flux.params(model)
loss(x,y) = Flux.mse(model(x), y)
# training
n_epochs = 1000
for epoch in 1:n_epochs
Flux.train!(loss, ps, training_set, opt)
println(sum([loss(i[1],i[2]) for i in training_set]))
end