Implement standard CIFAR-100 model in fedjax.models.cifar100
- 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
Assessment
This issue has not been assessed yet.