Partial parameter optimization in NNX (Transfer Learning)

[cgarciae]

Here is a recipe for performing partial optimization of a model in NNX using Optimizer and nnx.grad. The main trick is to define the differentiable state using the wrt filter for Optimizer, and passing that same filter to DiffState which is used to specify the gradient in nnx.grad.

class Classifier(nnx.Module):
  def __init__(self, embed_dim, num_classes, backbone, rngs):
    self.backbone = backbone
    self.head = nnx.Linear(embed_dim, num_classes, rngs=rngs)

  def __call__(self, x):
    x = self.backbone(x)
    x = self.head(x)
    return x

def load_model():
  return nnx.Linear(784, 1024, rngs=nnx.Rngs(0))

backbone = load_model()
classifier = Classifier(1024, 10, backbone, rngs=nnx.Rngs(1))

# filter to select only Params on head path
head_params = nnx.All(nnx.Param, nnx.PathContains('head'))

optimizer = nnx.Optimizer(
  classifier,
  tx=optax.adamw(3e-4),
  wrt=head_params,  # filter head params
)

# simple train step
@nnx.jit
def train_step(model, optimizer, x, y):
  def loss_fn(model):
    logits = model(x)
    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()

  diff_state = nnx.DiffState(0, head_params) # filter head params of the first argument
  grads = nnx.grad(loss_fn, argnums=diff_state)(model)
  optimizer.update(grads)

x = jnp.ones((1, 784))
y = jnp.ones((1,), jnp.int32)
train_step(classifier, optimizer, x, y)

[carlesoctav]

nice api :D