# Ensembling on multiple devices#

We show how to train an ensemble of CNNs on the MNIST dataset, where the size of the ensemble is equal to the number of available devices. In short, this change be described as:

make a number of functions parallel using

`jax.pmap()`

,split the random seed to obtain different parameter initialization,

replicate the inputs and unreplicate the outputs where necessary,

average probabilities across devices to compute the predictions.

In this HOWTO we omit some of the code such as imports, the CNN module, and metrics computation, but they can be found in the MNIST example.

## Parallel functions#

We start by creating a parallel version of `create_train_state()`

, which
retrieves the initial parameters of the models. We do this using `jax.pmap()`

.
The effect of “pmapping” a function is that it will compile the function with
XLA (similar to `jax.jit()`

), but execute it in parallel on XLA devices (e.g.,
GPUs/TPUs).

```
def create_train_state(rng, learning_rate, momentum):
cnn = CNN()
params = cnn.init(rng, jnp.ones([1, 28, 28, 1]))['params']
tx = optax.sgd(learning_rate, momentum)
return train_state.TrainState.create(
apply_fn=cnn.apply, params=params, tx=tx)
```

```
@functools.partial(jax.pmap, static_broadcasted_argnums=(1, 2))
def create_train_state(rng, learning_rate, momentum):
cnn = CNN()
params = cnn.init(rng, jnp.ones([1, 28, 28, 1]))['params']
tx = optax.sgd(learning_rate, momentum)
return train_state.TrainState.create(
apply_fn=cnn.apply, params=params, tx=tx)
```

Note that for the single-model code above, we use `jax.jit()`

to lazily
initialize the model (see Module.init’s documentation for more details).
For the ensembling case, `jax.pmap()`

will map over the first axis of the
provided argument `rng`

by default, so we should make sure that we provide
a different value for each device when we call this function later on.

Note also how we specify that `learning_rate`

and `momentum`

are static
arguments, which means the concrete values of these arguments will be used,
rather than abstract shapes. This is necessary because the provided arguments
will be scalar values. For more details see JIT mechanics: tracing and static
variables.

Next we simply do the same for the functions `apply_model()`

and
`update_model()`

. To compute the predictions from the ensemble, we take the
average of the individual probabilities. We use `jax.lax.pmean()`

to compute
the average *across devices*. This also requires us to specify the
`axis_name`

to both `jax.pmap()`

and `jax.lax.pmean()`

.

```
@jax.jit
def apply_model(state, images, labels):
def loss_fn(params):
logits = CNN().apply({'params': params}, images)
one_hot = jax.nn.one_hot(labels, 10)
loss = optax.softmax_cross_entropy(logits=logits, labels=one_hot).mean()
return loss, logits
grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
(loss, logits), grads = grad_fn(state.params)
accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
return grads, loss, accuracy
@jax.jit
def update_model(state, grads):
return state.apply_gradients(grads=grads)
```

```
@functools.partial(jax.pmap, axis_name='ensemble')
def apply_model(state, images, labels):
def loss_fn(params):
logits = CNN().apply({'params': params}, images)
one_hot = jax.nn.one_hot(labels, 10)
loss = optax.softmax_cross_entropy(logits=logits, labels=one_hot).mean()
return loss, logits
grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
(loss, logits), grads = grad_fn(state.params)
probs = jax.lax.pmean(jax.nn.softmax(logits), axis_name='ensemble')
accuracy = jnp.mean(jnp.argmax(probs, -1) == labels)
return grads, loss, accuracy
@jax.pmap
def update_model(state, grads):
return state.apply_gradients(grads=grads)
```

## Training the Ensemble#

Next we transform the `train_epoch()`

function. When calling the pmapped
functions from above, we mainly need to take care of duplicating the arguments
for all devices where necessary, and de-duplicating the return values.

```
def train_epoch(state, train_ds, batch_size, rng):
train_ds_size = len(train_ds['image'])
steps_per_epoch = train_ds_size // batch_size
perms = jax.random.permutation(rng, len(train_ds['image']))
perms = perms[:steps_per_epoch * batch_size]
perms = perms.reshape((steps_per_epoch, batch_size))
epoch_loss = []
epoch_accuracy = []
for perm in perms:
batch_images = train_ds['image'][perm, ...]
batch_labels = train_ds['label'][perm, ...]
grads, loss, accuracy = apply_model(state, batch_images, batch_labels)
state = update_model(state, grads)
epoch_loss.append(loss)
epoch_accuracy.append(accuracy)
train_loss = np.mean(epoch_loss)
train_accuracy = np.mean(epoch_accuracy)
return state, train_loss, train_accuracy
```

```
def train_epoch(state, train_ds, batch_size, rng):
train_ds_size = len(train_ds['image'])
steps_per_epoch = train_ds_size // batch_size
perms = jax.random.permutation(rng, len(train_ds['image']))
perms = perms[:steps_per_epoch * batch_size]
perms = perms.reshape((steps_per_epoch, batch_size))
epoch_loss = []
epoch_accuracy = []
for perm in perms:
batch_images = jax_utils.replicate(train_ds['image'][perm, ...])
batch_labels = jax_utils.replicate(train_ds['label'][perm, ...])
grads, loss, accuracy = apply_model(state, batch_images, batch_labels)
state = update_model(state, grads)
epoch_loss.append(jax_utils.unreplicate(loss))
epoch_accuracy.append(jax_utils.unreplicate(accuracy))
train_loss = np.mean(epoch_loss)
train_accuracy = np.mean(epoch_accuracy)
return state, train_loss, train_accuracy
```

As can be seen, we do not have to make any changes to the logic around the
`state`

. This is because, as we will see below in our training code,
the train state is replicated already, so when we pass it to `train_step()`

,
things will just work fine since `train_step()`

is pmapped. However,
the train dataset is not yet replicated, so we do that here. Since replicating
the entire train dataset is too memory intensive we do it at the batch level.

We can now rewrite the actual training logic. This consists of two simple
changes: making sure the RNGs are replicated when we pass them to
`create_train_state()`

, and replicating the test dataset, which is much
smaller than the train dataset so we can do this for the entire dataset
directly.

```
train_ds, test_ds = get_datasets()
rng = jax.random.key(0)
rng, init_rng = jax.random.split(rng)
state = create_train_state(init_rng, learning_rate, momentum)
for epoch in range(1, num_epochs + 1):
rng, input_rng = jax.random.split(rng)
state, train_loss, train_accuracy = train_epoch(
state, train_ds, batch_size, input_rng)
_, test_loss, test_accuracy = apply_model(
state, test_ds['image'], test_ds['label'])
logging.info(
'epoch:% 3d, train_loss: %.4f, train_accuracy: %.2f, '
'test_loss: %.4f, test_accuracy: %.2f'
% (epoch, train_loss, train_accuracy * 100, test_loss,
test_accuracy * 100))
```

```
train_ds, test_ds = get_datasets()
test_ds = jax_utils.replicate(test_ds)
rng = jax.random.key(0)
rng, init_rng = jax.random.split(rng)
state = create_train_state(jax.random.split(init_rng, jax.device_count()),
learning_rate, momentum)
for epoch in range(1, num_epochs + 1):
rng, input_rng = jax.random.split(rng)
state, train_loss, train_accuracy = train_epoch(
state, train_ds, batch_size, input_rng)
_, test_loss, test_accuracy = jax_utils.unreplicate(
apply_model(state, test_ds['image'], test_ds['label']))
logging.info(
'epoch:% 3d, train_loss: %.4f, train_accuracy: %.2f, '
'test_loss: %.4f, test_accuracy: %.2f'
% (epoch, train_loss, train_accuracy * 100, test_loss,
test_accuracy * 100))
```