# Working with models in FedJAX

In this chapter, we will learn about `fedjax.Model`

. This notebook assumes you already have finished the “Datasets” chapter. We first overview centralized training and evaluation with `fedjax.Model`

and then describe how to add new neural architectures and specify additional evaluation metrics.

```
# Uncomment these to install fedjax.
# !pip install fedjax
# !pip install --upgrade git+https://github.com/google/fedjax.git
```

```
import itertools
import jax
import jax.numpy as jnp
from jax.example_libraries import stax
import fedjax
```

## Centralized training & evaluation with `fedjax.Model`

Most federated learning algorithms are built upon common components from standard centralized learning. `fedjax.Model`

holds these common components. In centralized learning, we are mostly concerned with two tasks:

Training: We want to optimize our model parameters on the training dataset.

Evaluation: We want to know the values of evaluation metrics (e.g. accuracy) of the current model parameters on a test dataset.

Let’s first see how we can carry out these two tasks on the EMNIST dataset with `fedjax.Model`

.

```
# Load train/test splits of the EMNIST dataset.
train, test = fedjax.datasets.emnist.load_data()
# As a start, let's simply use a logistic regression model.
model = fedjax.models.emnist.create_logistic_model()
```

### Random initialization, the JAX way

To start training, we need some randomly initialized parameters. In JAX, pseudo random number generation works slightly differently. For now, it is sufficient to know we call `jax.random.PRNGKey()`

to seed the random number generator. JAX has a detailed introduction on this topic, if you are interested.

To create the initial model parameters, we simply call `fedjax.Model.init()`

with a `PRNGKey`

.

```
params_rng = jax.random.PRNGKey(0)
params = model.init(params_rng)
```

```
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
```

Here are our initial model parameters. With the same `PRNGKey`

, we will always get the same random initialization. There are 2 parameters in our model, the weights `w`

, and the bias `b`

. They are organized into a `FlapMapping`

, but in general any PyTree can be used to store model parameters.

```
params
```

```
FlatMapping({
'linear': FlatMapping({
'b': DeviceArray([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0.], dtype=float32),
'w': DeviceArray([[-0.04067196, 0.02348138, -0.0214883 , ..., 0.01055492,
-0.06988288, -0.02952586],
[-0.03985253, -0.03804361, 0.01401524, ..., 0.02281437,
-0.01771905, 0.06676884],
[ 0.00098182, -0.00844628, 0.01303554, ..., -0.05299249,
0.01777634, -0.0006488 ],
...,
[-0.05691862, 0.05192501, 0.01588603, ..., 0.0157204 ,
-0.01854135, 0.00297953],
[ 0.01680706, 0.05579231, 0.0459589 , ..., 0.01990358,
-0.01944044, -0.01710149],
[-0.00880739, 0.04229043, 0.00998938, ..., -0.00633441,
-0.04824542, 0.01395545]], dtype=float32),
}),
})
```

### Evaluating model parameters

Before we start training, let’s first see how our initial parameters fare on the train and test sets. Unsurprisingly, they do not do very well. We evaluate using the `fedjax.evaluate_model()`

which takes in model, parameters, and datasets which are batched. As noted in the dataset tutorial, we batch using
`fedjax.padded_batch_federated_data()`

for efficiency. `fedjax.padded_batch_federated_data()`

is very similar to `fedjax.ClientDataset.padded_batch()`

but operates over the entire federated dataset.

```
# We select first 16 batches using itertools.islice.
batched_test_data = list(itertools.islice(
fedjax.padded_batch_federated_data(test, batch_size=128), 16))
batched_train_data = list(itertools.islice(
fedjax.padded_batch_federated_data(train, batch_size=128), 16))
print('eval_test', fedjax.evaluate_model(model, params, batched_test_data))
print('eval_train', fedjax.evaluate_model(model, params, batched_train_data))
```

```
eval_test {'accuracy': DeviceArray(0.01757812, dtype=float32), 'loss': DeviceArray(4.1253214, dtype=float32)}
eval_train {'accuracy': DeviceArray(0.02490234, dtype=float32), 'loss': DeviceArray(4.116228, dtype=float32)}
```

How does our model know what evaluation metrics to report? It is simply specified in the `eval_metrics`

field. We will discuss evaluation metrics in more detail later.

```
model.eval_metrics
```

```
{'accuracy': Accuracy(target_key='y', pred_key=None),
'loss': CrossEntropyLoss(target_key='y', pred_key=None)}
```

Since `fedjax.evaluate_model()`

simply takes a stream of batches, we can also use it to evaluate multiple clients.

```
for client_id, dataset in itertools.islice(test.clients(), 4):
print(
client_id,
fedjax.evaluate_model(model, params,
dataset.padded_batch(batch_size=128)))
```

```
b'002d084c082b8586:f0185_23' {'accuracy': DeviceArray(0.05, dtype=float32), 'loss': DeviceArray(4.1247168, dtype=float32)}
b'005fdad281234bc0:f0151_02' {'accuracy': DeviceArray(0.09375, dtype=float32), 'loss': DeviceArray(4.093891, dtype=float32)}
b'014c177da5b15a39:f1565_04' {'accuracy': DeviceArray(0., dtype=float32), 'loss': DeviceArray(4.127692, dtype=float32)}
b'0156df0c34a25944:f3772_10' {'accuracy': DeviceArray(0.05263158, dtype=float32), 'loss': DeviceArray(4.1521378, dtype=float32)}
```

### The training objective

To train our model, we need two things: the objective function to minimize and an optimizer.

`fedjax.Model`

contains two functions that can be used to arrive at the training objective:

`apply_for_train(params, batch_example, rng)`

takes the current model parameters, a batch of examples, and a`PRNGKey`

, and returns some output.`train_loss(batch_example, train_output)`

translates the output of`apply_for_train()`

into a vector of per-example loss values.

In our example model, `apply_for_train()`

produces a score for each class and `train_loss()`

is simply the cross entropy loss. `apply_for_train()`

in this case does not make use of a `PRNGKey`

, so we can pass `None`

instead for convenience. A different `apply_for_train()`

might actually make use of the `PRNGKey`

, for tasks such as dropout.

```
# train_batches is an infinite stream of shuffled batches of examples.
def train_batches():
return fedjax.shuffle_repeat_batch_federated_data(
train,
batch_size=8,
client_buffer_size=16,
example_buffer_size=1024,
seed=0)
# We obtain the first batch by using the `next` function.
example = next(train_batches())
output = model.apply_for_train(params, example, None)
per_example_loss = model.train_loss(example, output)
output.shape, per_example_loss
```

```
((8, 62), DeviceArray([4.0337796, 4.046219 , 3.9447758, 3.933005 , 4.116893 ,
4.209843 , 4.060939 , 4.19899 ], dtype=float32))
```

Note that the `output`

is per example predictions and has shape (8, 62), where 8 is the batch size and 62 is the number of classes. Alternatively, we can use `model_per_example_loss()`

to get a function that gives us the same result. `model_per_example_loss()`

is a convenience function that does exactly what we just did.

```
per_example_loss_fn = fedjax.model_per_example_loss(model)
per_example_loss_fn(params, example, None)
```

```
DeviceArray([4.0337796, 4.046219 , 3.9447758, 3.933005 , 4.116893 ,
4.209843 , 4.060939 , 4.19899 ], dtype=float32)
```

The training objective is a scalar, so why does `train_loss()`

return a vector of per-example loss values? First of all, the training objective in most cases is just the average of the per-example loss values, so arriving at the final training objective isn’t hard. Moreover, in certain algorithms, we not only use the train loss over a single batch of examples for a stochastic training step, but also need to estimate the average train loss over an entire (client) dataset. Having the per-example loss values there is instrumental in obtaining the correct estimate when the batch sizes may vary.

```
def train_objective(params, example):
return jnp.mean(per_example_loss_fn(params, example, None))
train_objective(params, example)
```

```
DeviceArray(4.0680556, dtype=float32)
```

### Optimizers

With the training objective at hand, we just need an optimizer to find some good model parameters that minimize it.

There are many optimizer implementations in JAX out there, but FedJAX doesn’t force one choice over any other. Instead, FedJAX provides a simple `fedjax.optimizers.Optimizer`

interface so a new optimizer implementation can be wrapped. For convenience, FedJAX provides some common optimizers wrapped from optax.

```
optimizer = fedjax.optimizers.adam(1e-3)
```

An optimizer is simply a pair of two functions:

`init(params)`

returns the initial optimizer state, such as initial values for accumulators of gradients.`apply(grads, opt_state, params)`

applies the gradients to update the current optimizer state and model parameters.

Instead of modifying `opt_state`

or `params`

, `apply()`

returns a new pair of optimizer state and model parameters. In JAX, it is common to express computations in this stateless/mutation free style, often referred to as functional programming, or pure functions. The pureness of functions is crucial to many features in JAX, so it is always good practice to write functions that do not modify its inputs. You have probably also noticed that all the functions of `fedjax.Model`

we have seen so far do not modify the model object itself (for example, `init()`

returns model parameters instead of setting some attribute of `model`

; `apply_for_train()`

takes model parameters as an input argument, instead of getting it from `model`

). FedJAX does this to keep all functions pure.

However, in the top level training loop, it is fine to mutate states since we are not in a function that may be transformed by JAX. Let’s run our first training step, which resulted in a slight decrease in objective on the same batch of examples.

To obtain the gradients, we use `jax.grad()`

which returns the gradient function. More details about `jax.grad()`

can be found from the JAX documentation.

```
opt_state = optimizer.init(params)
grads = jax.grad(train_objective)(params, example)
opt_state, params = optimizer.apply(grads, opt_state, params)
train_objective(params, example)
```

```
DeviceArray(4.0080366, dtype=float32)
```

Instead of using `jax.grad()`

directly, FedJAX also provides a convenient `fedjax.model_grad()`

which computes the gradient of a model with respect to the averaged `fedjax.model_per_example_loss()`

.

```
model_grads = fedjax.model_grad(model)(params, example, None)
opt_state, params = optimizer.apply(grads, opt_state, params)
train_objective(params, example)
```

```
DeviceArray(3.9482572, dtype=float32)
```

Let’s wrap everything into a single JIT compiled function and train a few more steps, and evaluate again.

```
@jax.jit
def train_step(example, opt_state, params):
grads = jax.grad(train_objective)(params, example)
return optimizer.apply(grads, opt_state, params)
for example in itertools.islice(train_batches(), 5000):
opt_state, params = train_step(example, opt_state, params)
print('eval_test', fedjax.evaluate_model(model, params, batched_test_data))
print('eval_train', fedjax.evaluate_model(model, params, batched_train_data))
```

```
eval_test {'accuracy': DeviceArray(0.6152344, dtype=float32), 'loss': DeviceArray(1.5562292, dtype=float32)}
eval_train {'accuracy': DeviceArray(0.59765625, dtype=float32), 'loss': DeviceArray(1.6278805, dtype=float32)}
```

## Building a custom model

`fedjax.Model`

was designed with customization in mind. We have already seen how to switch to a different training loss. In this section, we will discuss how the rest of a `fedjax.Model`

can be customized.

### Training loss

Because `train_loss()`

is separate from `apply_for_train()`

, it is easy to switch to a different loss function.

```
def hinge_loss(example, output):
label = example['y']
num_classes = output.shape[-1]
mask = jax.nn.one_hot(label, num_classes)
label_score = jnp.sum(output * mask, axis=-1)
best_score = jnp.max(output + 1 - mask, axis=-1)
return best_score - label_score
hinge_model = model.replace(train_loss=hinge_loss)
fedjax.model_per_example_loss(hinge_model)(params, example, None)
```

```
DeviceArray([4.306656 , 0. , 0. , 0.4375435 , 0.96986485,
0. , 0.3052401 , 1.3918507 ], dtype=float32)
```

### Evaluation metrics

We have already seen that the `eval_metrics`

field of a `fedjax.Model`

tells the model what metrics to evaluate. `eval_metrics`

is a mapping from metric names to `fedjax.metrics.Metric`

objects. A `fedjax.metrics.Metric`

object tells us how to calculate a metric’s value from multiple batches of examples. Like `fedjax.Model`

, a `fedjax.metrics.Metric`

is stateless.

To customize the metrics to evaluate on, or what names to give to each, simply specify a different mapping.

```
only_accuracy = model.replace(
eval_metrics={'accuracy': fedjax.metrics.Accuracy()})
fedjax.evaluate_model(only_accuracy, params, batched_test_data)
```

```
{'accuracy': DeviceArray(0.6152344, dtype=float32)}
```

There are already some concrete `Metric`

s in `fedjax.metrics`

. It is also easy to implement a new one. You can read more about how to implement a `Metric`

in its own introduction.

The bit of `fedjax.Model`

that is directly relevant to evaluation is `apply_for_eval()`

. The relation between `apply_for_eval()`

and an evaluation metric is similar to that between `apply_for_train()`

and `train_loss()`

: `apply_for_eval(params, example)`

takes the model parameters and a batch of examples (notice there is no randomness in evaluation so we don’t need a `PRNGKey`

), and produces some prediction that evaluation metrics can consume. In our example, the outputs from `apply_for_eval()`

and `apply_for_train()`

are identical, but they don’t have to be.

```
jnp.all(
model.apply_for_train(params, example, None) == model.apply_for_eval(
params, example))
```

```
DeviceArray(True, dtype=bool)
```

What `apply_for_eval()`

needs to produce really just depends on what evaluation `fedjax.metrics.Metric`

s will be used. In our case, we are using `fedjax.metrics.Accuracy`

, and `fedjax.metrics.CrossEntropyLoss`

. They are similar in their requirements on the inputs:

They both need to know the true label from the

`example`

, using a`target_key`

that defaults to`"y"`

.They both need to know the predicted scores from

`apply_for_eval()`

, customizable as`pred_key`

. If`pred_key`

is None,`apply_for_eval()`

should return just a vector of per-class scores; otherwise`pred_key`

can be a string key, and`apply_for_eval()`

should return a mapping (e.g.`dict`

) that maps the key to a vector of per-class scores.

```
fedjax.metrics.Accuracy()
```

```
Accuracy(target_key='y', pred_key=None)
```

### Neural network architectures

We have now covered all five parts of a `fedjax.Model`

, namely `init()`

, `apply_for_train()`

, `apply_for_eval()`

, `train_loss()`

, and `eval_metrics`

. `train_loss()`

and `eval_metrics`

are easy to customize since they are mostly agnostic to the actual neural network architecture of the model. `init()`

, `apply_for_train()`

, and `apply_for_eval()`

on the other hand, are closely related.

In principle, as long as these three functions meet the interface we have seen so far, they can be used to build a custom model. Let’s try to build a model that uses multi-layer perceptron and hinge loss.

```
def cross_entropy_loss(example, output):
label = example['y']
num_classes = output.shape[-1]
mask = jax.nn.one_hot(label, num_classes)
return -jnp.sum(jax.nn.log_softmax(output) * mask, axis=-1)
def mlp_model(num_input_units, num_units, num_classes):
def mlp_init(rng):
w0_rng, w1_rng = jax.random.split(rng)
w0 = jax.random.uniform(w0_rng, [num_input_units, num_units])
b0 = jnp.zeros([num_units])
w1 = jax.random.uniform(w1_rng, [num_units, num_classes])
b1 = jnp.zeros([num_classes])
return w0, b0, w1, b1
def mlp_apply(params, batch, rng=None):
w0, b0, w1, b1 = params
x = batch['x']
batch_size = x.shape[0]
h = jax.nn.relu(x.reshape([batch_size, -1]) @ w0 + b0)
return h @ w1 + b1
return fedjax.Model(
init=mlp_init,
apply_for_train=mlp_apply,
apply_for_eval=mlp_apply,
train_loss=cross_entropy_loss,
eval_metrics={'accuracy': fedjax.metrics.Accuracy()})
# There are 28*28 input pixels, and 62 classes in EMNIST.
mlp = mlp_model(28 * 28, 128, 62)
@jax.jit
def mlp_train_step(example, opt_state, params):
@jax.grad
def grad_fn(params, example):
return jnp.mean(fedjax.model_per_example_loss(mlp)(params, example, None))
grads = grad_fn(params, example)
return optimizer.apply(grads, opt_state, params)
params = mlp.init(jax.random.PRNGKey(0))
opt_state = optimizer.init(params)
print('eval_test before training:',
fedjax.evaluate_model(mlp, params, batched_test_data))
for example in itertools.islice(train_batches(), 5000):
opt_state, params = mlp_train_step(example, opt_state, params)
print('eval_test after training:',
fedjax.evaluate_model(mlp, params, batched_test_data))
```

```
eval_test before training: {'accuracy': DeviceArray(0.05078125, dtype=float32)}
eval_test after training: {'accuracy': DeviceArray(0.4951172, dtype=float32)}
```

While writing custom neural network architectures from scratch is possible, most of the time, it is much more convenient to use a neural network library such as Haiku or `jax.example_libraries.stax`

. The two functions `fedjax.create_model_from_haiku`

and `fedjax.create_model_from_stax`

can convert a neural network expressed in the respective framework into a `fedjax.Model`

. Let’s build a convolutional network using `jax.example_libraries.stax`

this time.

```
def stax_cnn_model(input_shape, num_classes):
stax_init, stax_apply = stax.serial(
stax.Conv(
out_chan=64, filter_shape=(3, 3), strides=(1, 1), padding='SAME'),
stax.Relu,
stax.Flatten,
stax.Dense(256),
stax.Relu,
stax.Dense(num_classes),
)
return fedjax.create_model_from_stax(
stax_init=stax_init,
stax_apply=stax_apply,
sample_shape=input_shape,
train_loss=cross_entropy_loss,
eval_metrics={'accuracy': fedjax.metrics.Accuracy()})
stax_cnn = stax_cnn_model([-1, 28, 28, 1], 62)
@jax.jit
def stax_cnn_train_step(example, opt_state, params):
@jax.grad
def grad_fn(params, example):
return jnp.mean(
fedjax.model_per_example_loss(stax_cnn)(params, example, None))
grads = grad_fn(params, example)
return optimizer.apply(grads, opt_state, params)
params = stax_cnn.init(jax.random.PRNGKey(0))
opt_state = optimizer.init(params)
print('eval_test before training:',
fedjax.evaluate_model(stax_cnn, params, batched_test_data))
for example in itertools.islice(train_batches(), 1000):
opt_state, params = stax_cnn_train_step(example, opt_state, params)
print('eval_test after training:',
fedjax.evaluate_model(stax_cnn, params, batched_test_data))
```

```
eval_test before training: {'accuracy': DeviceArray(0.03076172, dtype=float32)}
eval_test after training: {'accuracy': DeviceArray(0.72558594, dtype=float32)}
```

## Recap

In this chapter, we have covered the following:

Components of

`fedjax.Model`

:`init()`

,`apply_for_train()`

,`apply_for_eval()`

,`train_loss()`

, and`eval_metrics`

.Optimizers in

`fedjax.optimizers`

.Standard centralized learning with a

`fedjax.Model`

.Specifying evaluation metrics in

`eval_metrics`

.Building a custom

`fedjax.Model`

.