Skip to content
Bayesian Nonparametric Federated Learning of Neural Networks
Branch: master
Clone or download
Type Name Latest commit message Commit time
Failed to load latest commit information.
data Commit PFNM code May 12, 2019
logs Commit PFNM code May 12, 2019
LICENSE Initial commit May 9, 2019 Update README with paper link May 28, 2019 Commit PFNM code May 12, 2019 Update dataset load semantics to new torchvision format May 29, 2019 Commit PFNM code May 12, 2019

Probabilistic Federated Neural Matching

This is the code accompanying the ICML 2019 paper "Bayesian Nonparametric Federated Learning of Neural Networks" Paper link: []

Requirements to run the code:

  1. Python 3.6
  2. PyTorch 0.4
  3. Scikit-learn
  4. Matplotlib
  5. Numpy

Important source files:

  1. Main entryway to the code. Used to run all experiments
  2. matching/ Contains the PFNM matching code for single communication federated learning
  3. matching/ Contains the PFNM matching code for multiple communication federated learning

Sample Commands:

  1. MNIST Heterogenous 10 batches

python --logdir "logs/mnist_test" --dataset "mnist" --datadir "data/mnist/" --net_config "784, 100, 10" --n_nets 10 --partition "hetero-dir" --experiment "u-ensemble,pdm,pdm_iterative" --lr 0.01 --epochs 10 --reg 1e-6 --communication_rounds 5 --lr_decay 0.99 --iter_epochs 5

  1. CIFAR-10 Heterogenous 10 batches

python --logdir "logs/cifar10_test" --dataset "cifar10" --datadir "data/cifar10/" --net_config "3072, 100, 10" --n_nets 10 --partition "hetero-dir" --experiment "u-ensemble,pdm,pdm_iterative" --lr 0.001 --epochs 10 --reg 1e-5 --communication_rounds 5 --lr_decay 0.99 --iter_epochs 5

Important arguments:

The following arguments to the PFNM file control the important parameters of the experiment

  1. net_config: Defines the local network architecture. CSV of sizes. Ex: "784, 100, 100, 10" defines a 2-layer network with 100 neurons in each layer.
  2. n_nets: Number of local networks. This is denoted by "J" in the paper
  3. partition: Kind of data partition. Values: homo, hetero-dir
  4. experiments: Defines which experiments will be executed. Values: u-ensemble (Uniform ensemble), pdm (PFNM matching), pdm_iterative (PFNM with extra communications)
  5. communication_rounds: How many rounds of communication between the local learner and the master network in the case of PFNM with multiple communications.


Some of the output is printed on the terminal. However, majority of the information is logged to a log file in the specified log folder.

Citing PFNM

  title = 	 {{B}ayesian Nonparametric Federated Learning of Neural Networks},
  author = 	 {Yurochkin, Mikhail and Agarwal, Mayank and Ghosh, Soumya and Greenewald, Kristjan and Hoang, Nghia and Khazaeni, Yasaman},
  booktitle = 	 {Proceedings of the 36th International Conference on Machine Learning},
  pages = 	 {7252--7261},
  year = 	 {2019},
  editor = 	 {Chaudhuri, Kamalika and Salakhutdinov, Ruslan},
  volume = 	 {97},
  series = 	 {Proceedings of Machine Learning Research},
  address = 	 {Long Beach, California, USA},
  month = 	 {09--15 Jun},
  publisher = 	 {PMLR},
  pdf = 	 {},
  url = 	 {},
  abstract = 	 {In federated learning problems, data is scattered across different servers and exchanging or pooling it is often impractical or prohibited. We develop a Bayesian nonparametric framework for federated learning with neural networks. Each data server is assumed to provide local neural network weights, which are modeled through our framework. We then develop an inference approach that allows us to synthesize a more expressive global network without additional supervision, data pooling and with as few as a single communication round. We then demonstrate the efficacy of our approach on federated learning problems simulated from two popular image classification datasets.}
You can’t perform that action at this time.