How to track the current iteration number in a custom loss function?

[GFODK]

Hi,
I'm training a symbolic regression model in Julia (not using PySR) and defining a custom loss function. My loss function includes two additional conditions, and I manually assign different weights to each term, like this:

w_a, w_b, w_c = Ref(0.1), Ref(0.4), Ref(0.4)

function custom_loss(tree::Node, dataset::Dataset{T,L}, options::Options) where {T,L}
    ...
    loss = L_a * w_a + L_b * w_b + L_c * w_c
    return loss
end

model = SRRegressor(
    niterations=1000,
    binary_operators=[+, -, *, /],
    unary_operators=[sin, cos, exp, log],
    population_size=50,
    loss_function=custom_loss
)

mach = machine(model, x, y)
fit!(mach)

I’d like to adjust the weights of my loss terms dynamically based on the current iteration.
Is there a built-in variable that keeps track of the iteration count inside the custom_loss function? Or is there another way to access the current iteration during training?

Any help would be much appreciated! Thanks! 😊

[MilesCranmer]

Hi @GFODK,

Unfortunately you cannot have a dynamic loss function, because it is cached at various points in the search rather than re-computing each time. So if your loss function changes, it will break some assumptions in the code.

What you can do though is adjust the loss function and call fit! again—it will start where it left off, and recompute all the losses on existing populations of expressions. So, for example:

model = SRRegressor(
    niterations=1,
    binary_operators=[+, -, *, /],
    unary_operators=[sin, cos, exp, log],
    population_size=50,
    loss_function=loss_functions[1]
)
mach = machine(model, x, y)
fit!(mach)

for i in 2:100
    mach.model.loss_function = loss_functions[i]
    mach.model.niterations += 1
    fit!(mach)
end

or something like this. You might want to build some callable struct that has parameters you can declare, like

struct LossContext <: Function
    w_a::Float64
    w_b::Float64
    w_c::Float64
end

function (ctx::LossContext)(tree::Node, dataset::Dataset{T,L}, options::Options) where {T,L}
    ...
    return L_a * ctx.w_a + L_b * ctx.w_b + L_c * ctx.w_c
end

and then you can just define the loss like

    SRRegressor(loss_function=LossContext(0.1, 0.4, 0.4))

[GFODK]

Hey, thanks so much for the quick response! 🙌

This is my first time working with Julia, so I hadn’t thought of that approach. I ended up rewriting my code by combining both of the methods you suggested. Now, I'm updating the loss weights (w_a, w_b, w_c) at each iteration and assigning them to loss_function.

However, I'm seeing the following warning messages during training. Do you think it's safe to ignore them, or could they indicate that something is going wrong with the training process?

Here’s the modified code—I'm updating the loss function in every iteration using a for loop:

struct LossContext <: Function
    w_a::Float64
    w_b::Float64
    w_c::Float64
end

w_a, w_b, w_c = 1.0, 1.0, 1.0
model = SRRegressor(
    niterations=1,
    binary_operators=[+, -, *, /],
    unary_operators=[sin, cos, exp, log],
    population_size=100,
    loss_function=LossContext(
        w_a=w_a, 
        w_b=w_b, 
        w_c=w_c
    )
)
mach = machine(model, x, y)
fit!(mach)

# Update weights for the next iteration
w_a, w_b, w_c = update_weights(w_a, w_b, w_c)

let w_a=w_a, w_b=w_b, w_c=w_c
    for i in 2:100
        mach.model.loss_function = LossContext(
            w_a=w_a, 
            w_b=w_b, 
            w_c=w_c
        )
        mach.model.niterations += 1
        fit!(mach)
        ...
        w_a, w_b, w_c = update_weights(w_a, w_b, w_c)
    end
end

And here are the warning messages I’m getting:

[ Info: Updating machine(SRRegressor(defaults = nothing, …), …).  
┌ Warning: Recreating population (output=1, population=2), as the saved one doesn't have the correct number of members.  
└ @ SymbolicRegression ~/.julia/packages/SymbolicRegression/kbyoS/src/SymbolicRegression.jl:743  
┌ Warning: Recreating population (output=1, population=16), as the saved one doesn't have the correct number of members.  
└ @ SymbolicRegression ~/.julia/packages/SymbolicRegression/kbyoS/src/SymbolicRegression.jl:743  
┌ Warning: Recreating population (output=1, population=18), as the saved one doesn't have the correct number of members.  
└ @ SymbolicRegression ~/.julia/packages/SymbolicRegression/kbyoS/src/SymbolicRegression.jl:743  
...

Would really appreciate any insights on whether this is expected behavior or if I should change something. Thanks again! 😊

[MilesCranmer]

Hm. I think it's fine. Maybe just increase by 2 (or more) iterations at a time instead of 1? Then it has a chance to repopulate before it needs to return the expressions