How to fine-tune neuralgcm?

[deokgue]

Is there any sample code for fine-tuning neuralgcm?

[shoyer]

PressureLevelModel objects are JAX pytrees, so you can differentiate with respect to them.

As a simple example, here's how you could minimize a reconstruction loss on the encoder/decoder part of a toy NeuralGCM model for a single time slice of data:

import jax
import optax
import neuralgcm

@jax.jit
def compute_loss(model, inputs, forcings, rng):
    encoded = model.encode(inputs, forcings, rng_key=rng)
    restored = model.decode(encoded, forcings)
    return abs(inputs['temperature'] - restored['temperature']).mean()

checkpoint = neuralgcm.demo.load_checkpoint_tl63_stochastic()
initial_model = neuralgcm.PressureLevelModel.from_checkpoint(checkpoint)

ds = neuralgcm.demo.load_data(initial_model.data_coords)
inputs, forcings = model.data_from_xarray(ds.isel(time=0))

print(compute_loss(initial_model, inputs, forcings, jax.random.key(0)))
# Array(6.2256584, dtype=float32)

optimizer = optax.adam(1e-3)

model = initial_model
opt_state = optimizer.init(model)

for i in range(5):
    loss, grads = jax.value_and_grad(compute_loss)(model, inputs, forcings, jax.random.key(0))
    updates, opt_state = optimizer.update(grads, opt_state)
    model = optax.apply_updates(model, updates)
    print(f'{i=}, {loss=}')
# i=0, loss=Array(6.2256584, dtype=float32)
# i=1, loss=Array(4.670498, dtype=float32)
# i=2, loss=Array(3.855668, dtype=float32)
# i=3, loss=Array(3.5485578, dtype=float32)
# i=4, loss=Array(3.4335625, dtype=float32)

This model quickly overfits to the single day's data (and the RNG key), but hopefully you get the idea.