Writing a Training Loop in JAX and Flax

# Initializing a Weights & Biases Run
wandb.init(
    project="simple-training-loop",
    entity="jax-series",
    job_type="simple-train-loop"
)

# Setting the configs of our experiment using `wandb.config`.
# This way, Weights & Biases automatcally syncs the configs of 
# our experiment which could be used to reproduce the results of an experiment.
config = wandb.config
config.seed = 42
config.batch_size = 64
config.validation_split = 0.2
config.pooling = "avg"
config.learning_rate = 1e-4
config.epochs = 15

class CNN(nn.Module):
    pool_module: Callable = nn.avg_pool

    def setup(self):
        self.conv_1 = nn.Conv(features=32, kernel_size=(3, 3))
        self.conv_2 = nn.Conv(features=32, kernel_size=(3, 3))
        self.conv_3 = nn.Conv(features=64, kernel_size=(3, 3))
        self.conv_4 = nn.Conv(features=64, kernel_size=(3, 3))
        self.conv_5 = nn.Conv(features=128, kernel_size=(3, 3))
        self.conv_6 = nn.Conv(features=128, kernel_size=(3, 3))
        self.dense_1 = nn.Dense(features=1024)
        self.dense_2 = nn.Dense(features=512)
        self.dense_output = nn.Dense(features=10)

    @nn.compact
    def __call__(self, x):
        x = nn.relu(self.conv_1(x))
        x = nn.relu(self.conv_2(x))
        x = self.pool_module(x, window_shape=(2, 2), strides=(2, 2))
        x = nn.relu(self.conv_3(x))
        x = nn.relu(self.conv_4(x))
        x = self.pool_module(x, window_shape=(2, 2), strides=(2, 2))
        x = nn.relu(self.conv_5(x))
        x = nn.relu(self.conv_6(x))
        x = self.pool_module(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))
        x = nn.relu(self.dense_1(x))
        x = nn.relu(self.dense_2(x))
        return self.dense_output(x)

rng = jax.random.PRNGKey(config.seed) # PRNG Key
x = jnp.ones(shape=(config.batch_size, 32, 32, 3)) # Dummy Input
model = CNN(pool_module=MODULE_DICT[config.pooling]) # Instantiate the Model
params = model.init(rng, x) # Initialize the parameters
jax.tree_map(lambda x: x.shape, params) # Check the parameters

import optax
from flax.training import train_state

def init_train_state(
    model, random_key, shape, learning_rate
) -> train_state.TrainState:
    # Initialize the Model
    variables = model.init(random_key, jnp.ones(shape))
    # Create the optimizer
    optimizer = optax.adam(learning_rate)
    # Create a State
    return train_state.TrainState.create(
        apply_fn = model.apply,
        tx=optimizer,
        params=variables['params']
    )


state = init_train_state(
    model, rng, (config.batch_size, 32, 32, 3), config.learning_rate
)

def train_and_evaluate(train_dataset, eval_dataset, test_dataset, state, epochs):
    num_train_batches = tf.data.experimental.cardinality(train_dataset)
    num_eval_batches = tf.data.experimental.cardinality(eval_dataset)
    num_test_batches = tf.data.experimental.cardinality(test_dataset)
   
    for epoch in tqdm(range(1, epochs + 1)):
        best_eval_loss = 1e6


        # ============== Training ============== #
        train_batch_metrics = []
        train_datagen = iter(tfds.as_numpy(train_dataset))
        for batch_idx in range(num_train_batches):
            batch = next(train_datagen)
            state, metrics = train_step(state, batch)
            train_batch_metrics.append(metrics)
        train_batch_metrics = accumulate_metrics(train_batch_metrics)


        # ============== Validation ============= #
        eval_batch_metrics = []
        eval_datagen = iter(tfds.as_numpy(eval_dataset))
        for batch_idx in range(num_eval_batches):
            batch = next(eval_datagen)
            metrics = eval_step(state, batch)
            eval_batch_metrics.append(metrics)
        eval_batch_metrics = accumulate_metrics(eval_batch_metrics)

        # Log Metrics to Weights & Biases
        wandb.log({
            "Train Loss": train_batch_metrics['loss'],
            "Train Accuracy": train_batch_metrics['accuracy'],
            "Validation Loss": eval_batch_metrics['loss'],
            "Validation Accuracy": eval_batch_metrics['accuracy']
        }, step=epoch)

    return state

@jax.jit
def train_step(
    state: train_state.TrainState, batch: jnp.ndarray
):
    image, label = batch

    def loss_fn(params):
        logits = state.apply_fn({'params': params}, image)
        loss = cross_entropy_loss(logits=logits, labels=label)
        return loss, logits

    gradient_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (_, logits), grads = gradient_fn(state.params)
    state = state.apply_gradients(grads=grads)
    metrics = compute_metrics(logits=logits, labels=label)
    return state, metrics

@jax.jit
def eval_step(state, batch):
    image, label = batch
    logits = state.apply_fn({'params': state.params}, image)
    return compute_metrics(logits=logits, labels=label)

def train_step(
    state: train_state.TrainState, batch: jnp.ndarray
):
    ...

# Now this jitted version of the train_step function would be used in the training loop
jitted_train_step = jax.jit(train_step)