Adding a new controller requires a lot of sub-changes. The easiest way is to extend the base class BaseController
and override the training hook method _train
. The base class takes care of initializing the experiment on all devices, starting the logging process and more. Generally speaking, your implementation wants to iterate over the number of rounds. For each round, you will have to:
- Do a training step
- Refresh the list of active devices in case some dropped out (using
_refresh_active_devices
) - Optionally: aggregate metrics and log them
- Optionally: save weights to file.
The best controller to take a look at is the FedController
. It has an easy implementation and shows how to call into the RPC layer. If your implementation is similar to an existing one, it is the best to start there and copy the old code into a new controller class and apply your modifications step by step. There are several things you might need or need not to change, depending on your use-case:
- Add new RPC calls (see here for a more detailed overview)
- Add new methods to the device class
- Optionally: add constructor arguments to configure the controller, if necessary
- Add a new config file under
edml/config
that showcases/defines a default configuration for the controller.