Custom loss functions in `Lux.jl`

jamblejoe
Fri, 07 Feb 2025 15:46:58 +0000
Machine Learning

I am currently getting accustomed to the Lux.jl package. I want to use a custom loss function which includes a l_2-regression term consisting of the model weights. I defined my custom loss function as

function loss_function(model, ps, st, (x, y))
    T = eltype(Base.Flatten(first(ps)))

    loss_mse = MSELoss()(model, ps, st, (x,y))[1]

    loss_reg = zero(T)
    for p in ps
        loss_reg += sum(abs2, Base.Flatten(p))
    end

    loss_total = loss_mse + convert(T, 0.001) * loss_reg

    return loss_total, st, NamedTuple()
end

I’d like to get some feedback of more experienced users of Lux.jl. Is this a good approach? It works with basic models in Lux.jl. My questions:

Is calling MSELoss() inside the function a good idea or better move it outside loss_function?

I use Enzyme.jl as AD backend, which likes type-stability. I struggled a bit to figure out the return type. Is my approach appropriate?

Thanks for the responses and those awesome packages!

6 posts - 2 participants

Read full topic