This is a proof-of-concept of federated learning using the Candle framework with Rust. It implements the FedAvg algorithm for horizontal federated learning with data provided by workers.
Multiple workers connect to a coordinator, which orchestrates them to train a model on their local data. The focus of this code is on the distributed system needed for federated learning, not on the machine learning model. As such, the model is a simple linear classification model and each worker trains on the same MNIST dataset.
The components communicate using gRPC in a client-server model.
A coordinator manages the training process. It provides a publish-subscribe service for workers to connect to, send training requests, and receive training results. These results are then aggregated by the coordinator. It also provides a service to start a training run.
Workers have access to their local training data. They connect to a coordinator and wait for training requests. When they receive a training request, they train a model on their local data and send the trained model back to the coordinator.
Build the project with cargo build -r and start the coordinator with
cargo run -r --bin coordinator then connect one or more workers with
cargo run -r --bin worker. With the workers connected, start a training run
with cargo run -r --bin start_training 10. This will train the model for 10
rounds. For example:
$ cargo run -r --bin coordinator &
$ cargo run -r --bin worker &
$ cargo run -r --bin worker &
$ cargo run -r --bin start_training 10