[Python] How to control hyperparameter within flax.nnx loss function using an optax.schedule?

from jax import numpy as jnp
from jax import random
from flax import nnx
import optax
from matplotlib import pyplot as plt

if __name__ == '__main__':
shape = (2,55,1)
epochs = 123
rngs = nnx.Rngs(123)
model = nnx.Linear( 1, 1, rngs=rngs )

skey = rngs.params()
xx = random.uniform( skey, shape, minval=-10, maxval=10 )

def f(x,m=2.234,b=-1.123):
return m*x+b
obs1,obs2 = f(xx)
x1,x2 = xx

c = 0.9
learning_rate_schedule = optax.schedules.cosine_decay_schedule(
init_value=2e-1,
decay_steps = int(c*epochs),
alpha= 0.01,
)
## how do I get the values from hyperparam_schedule into the loss function?
hyperparam_schedule = optax.schedules.linear_schedule(
init_value=12,
end_value=234,
transition_steps=int(c*epochs),
)

params = nnx.Param
optimizer = nnx.Optimizer(
model,
tx=optax.adam(learning_rate_schedule),
wrt = params
)

@nnx.scan(
in_axes=(nnx.Carry,None,None),
out_axes=(nnx.Carry,0),
length=epochs
)
def optimizer_scan( carry, x, obs ):
def loss_function(model, inputs, obs):
prediction = model(inputs)
## I want this to instead look like:
#prediction = model( my_hyperparameter, inputs )
error = obs - prediction
loss = jnp.mean(error ** 2)
mae = jnp.mean(jnp.abs(error ) )
return loss, mae

model, optimizer= carry
(loss,mae), grads = nnx.value_and_grad(loss_function,has_aux = True)( model, x, obs )
optimizer.update( model, grads )
return (model,optimizer), (loss,mae)
(model,optimizer), (losses,maes) = optimizer_scan( (model, optimizer), x1, obs1 )

print( ' AFTER TRAINING' )
print( 'training loss:', losses[-1] )

y1,y2 = model(xx)
error = obs2-y2
loss = jnp.mean( error*error )
print( 'test loss:',loss )
print( 'm approximation:', model.kernel.value )
print( 'b approximation:', model.bias.value )


I want to control an argument to my model's call method with an optax schedule. This is a hyperparameter, and I want it to change slowly over training. I have marked the schedule that I have already defined for this purpose with a comment. I have also shown in the loss function what I want the forward pass call to look like.

I see that there is an optax.schedules method to inject hyperparameters, but I couldn't get it to do something useful here.

PS This is updated for and tested on flax version 0.11

Continue reading...