Step by Step Guide to Build and Compare FedAvg and FedProx Federated Learning on Non-IID CIFAR-10 with NVIDIA FLARE

In this tutorial, we build an advanced federated learning experiment with NVIDIA FLARE. We compare FedAvg and FedProx on a non-IID CIFAR-10 setup, where client data is split using a Dirichlet distribution to simulate realistic label imbalance across federated sites. We use the NVFlare Job API to define and launch federated jobs, while the Client API handles local training, model exchange, and communication between simulated clients and the server. Finally, we run both algorithms on the same partitioned dataset and visualize how their global model accuracy evolves across communication rounds.

!pip install -q "nvflare>=2.5" torch torchvision matplotlib

import os, glob, shutil, sys

import numpy as np

import torch, torchvision