google / google/fedjax

Implement standard CIFAR-100 model in fedjax.models.cifar100

Open
#268 1 comment 0 reactions 0 assignees View on GitHub
contributions welcome enhancement
Dominant language
Python
Stars
272
Forks
42
PR merge metrics
No merged PRs in 30d

Description

Add a standard implementation of the model for the CIFAR-100 task. The dataset can be found in [fedjax.datasets.cifar100](https://github.com/google/fedjax/blob/main/fedjax/datasets/cifar100.py).

For the model architecture, we should follow [“Adaptive Federated Optimization”](https://arxiv.org/abs/2003.00295). The model architecture is detailed in section 4 as a ResNet-18 (replacing batch norm with group norm). Code for this paper and a Keras implementation of the model can be found [here](https://github.com/google-research/federated/blob/master/utils/models/resnet_models.py). We suggest using either [haiku](https://github.com/deepmind/dm-haiku) or [flax](https://github.com/google/flax) to implement the model for use with JAX.

If you choose to use haiku, you can use [fedjax.create_model_from_haiku](https://fedjax.readthedocs.io/en/stable/fedjax.html#fedjax.create_model_from_haiku) to create a fedjax compatible model. If you choose to use flax, wrapping it in a [fedjax.Model](https://fedjax.readthedocs.io/en/stable/fedjax.html#fedjax.Model) is fairly straightforward and we can provide guidance for this.

A good example to follow is #265 that checks in a simple linear model for CIFAR-100 and includes the model implementation, tests, and baseline results with FedAvg using [this script](https://github.com/google/fedjax/blob/main/experiments/fed_avg/run_fed_avg.py). Make sure to add a flags file similar to https://github.com/google/fedjax/blob/main/experiments/fed_avg/fed_avg.CIFAR100_LOGISTIC.flags and add the new task to https://github.com/google/fedjax/blob/main/fedjax/training/tasks.py.

Thanks for your contributions!

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.