fluke.nets

This module contains the definition of several neural networks used in state-of-the-art federated learning papers.

EncoderHeadNet	Encoder (aka backbone) + Head Network [Base Class] This type of networks are defined as two subnetworks, where one is meant to be the encoder/backbone network that learns a latent representation of the input, and the head network that is the classifier part of the model.
GlobalLocalNet	Global-Local Network (Abstract Class).
HeadGlobalEncoderLocalNet	This implementation of the Global-Local Network (GlobalLocalNet) is meant to be used with the Encoder-Head architecture.
EncoderGlobalHeadLocalNet	This implementation of the Global-Local Network (GlobalLocalNet) is meant to be used with the Encoder-Head architecture.
MNIST_2NN	Multi-layer Perceptron for MNIST.
MNIST_2NN_E	Encoder for the MNIST_2NN network.
MNIST_2NN_D	Head for the MNIST_2NN network.
MNIST_CNN	Convolutional Neural Network for MNIST.
MNIST_CNN_E	Encoder for the MNIST_CNN network.
MNIST_CNN_D	Head for the MNIST_CNN network.
FedBN_CNN	Convolutional Neural Network with Batch Normalization for CIFAR-10.
FedBN_CNN_E	Encoder for the FedBN_CNN network.
FedBN_CNN_D	Head for the FedBN_CNN network.
MNIST_LR	Logistic Regression for MNIST.
CifarConv2	Convolutional Neural Network for CIFAR-10.
CifarConv2_E	Encoder for the CifarConv2 network.
CifarConv2_D	Head for the CifarConv2 network.
ResNet9	ResNet-9 network for CIFAR-100 classification.
ResNet9_E	Encoder for the ResNet9 network.
ResNet9_D	Head for the ResNet9 network.
FEMNIST_CNN	Convolutional Neural Network for FEMNIST.
FEMNIST_CNN_E	Encoder for the FEMNIST_CNN network.
FEMNIST_CNN_D	Head for the FEMNIST_CNN network.
VGG9_E	Encoder for the VGG9 network.
VGG9_D	Head for the VGG9 network.
VGG9	VGG-9 network for FEMNIST classification.
ResNet18	ResNet-18 network as defined in the torchvision library.
ResNet34	ResNet-34 network as defined in the torchvision library.
ResNet50	ResNet-50 network as defined in the torchvision library.
ResNet18GN	ResNet-18 network as defined in the torchvision library but with Group Normalization layers instead of Batch Normalization.
MoonCNN_E	Encoder for the MoonCNN network.
MoonCNN_D	Head for the MoonCNN network.
MoonCNN	Convolutional Neural Network for CIFAR-10.
LeNet5_E	Encoder for the LeNet5 network.
LeNet5_D	Head for the LeNet5 network.
LeNet5	LeNet-5 for CIFAR.
Shakespeare_LSTM_E	Encoder for the Shakespeare_LSTM network.
Shakespeare_LSTM_D	Head for the Shakespeare_LSTM network.
Shakespeare_LSTM	LSTM for Shakespeare.

Classes

This module contains the definition of several neural networks used in state-of-the-art federated learning papers.

class fluke.nets.EncoderHeadNet(encoder: Module, head: Module)

Bases: Module

Encoder (aka backbone) + Head Network [Base Class] This type of networks are defined as two subnetworks, where one is meant to be the encoder/backbone network that learns a latent representation of the input, and the head network that is the classifier part of the model. The forward method should work as usual (i.e., \(g(f(\mathbf{x}))\) where \(\mathbf{x}\) is the input, \(f\) is the encoder and \(g\) is the head), but the forward_encoder and forward_head methods should be used to get the output of the encoder and head subnetworks, respectively. If this is not possible, they fallback to the forward method (default behavior).

Parameters:

• encoder (nn.Module) – Encoder subnetwork.

• head (nn.Module) – Head subnetwork.

property encoder: Module

Return the encoder subnetwork.

Returns:

Encoder subnetwork.

Return type:

nn.Module

property head: Module

Return the head subnetwork.

Returns:

head subnetwork.

Return type:

nn.Module

forward_encoder(x: Tensor) → Tensor

Forward pass through the encoder subnetwork.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor of the encoder subnetwork.

Return type:

torch.Tensor

forward_head(z: Tensor) → Tensor

Forward pass through the head subnetwork. z is assumed to be the output of the encoder subnetwork or an “equivalent” tensor.

Parameters:

z (torch.Tensor) – Input tensor.

Returns:

Output tensor of the head subnetwork.

Return type:

torch.Tensor

class fluke.nets.GlobalLocalNet(*args, **kwargs)

Bases: Module

Global-Local Network (Abstract Class). This is a network that has two subnetworks, one is meant to be shared (global) and one is meant to be personalized (local). The forward method should work as expected, but the forward_local and forward_global methods should be used to get the output of the local and global subnetworks, respectively. If this is not possible, they fallback to the forward method (default behavior).

abstract get_local() → Module

Return the local subnetwork.

Returns:

The local subnetwork

Return type:

nn.Module

abstract get_global() → Module

Return the global subnetwork.

Returns:

The global subnetwork

Return type:

nn.Module

class fluke.nets.HeadGlobalEncoderLocalNet(model: EncoderHeadNet)

Bases: GlobalLocalNet

This implementation of the Global-Local Network (GlobalLocalNet) is meant to be used with the Encoder-Head architecture. The global (i.e., that is shared between clients and server) subnetwork is the head and the local (i.e., not shared) subnetwork is the encoder.

Parameters:

model (EncoderHeadNet) – The federated model to use.

See also

• EncoderHeadNet

• GlobalLocalNet

• EncoderGlobalHeadLocalNet

class fluke.nets.EncoderGlobalHeadLocalNet(model: EncoderHeadNet)

Bases: GlobalLocalNet

This implementation of the Global-Local Network (GlobalLocalNet) is meant to be used with the Encoder-Head architecture. The global (i.e., that is shared between clients and server) subnetwork is the encoder and the local (i.e., not shared) subnetwork is the head.

Parameters:

model (EncoderHeadNet) – The federated model to use.

See also

• EncoderHeadNet

• GlobalLocalNet

• HeadGlobalEncoderLocalNet

class fluke.nets.MNIST_2NN(hidden_size: tuple[int, int] = (200, 100), softmax: bool = False)

Bases: EncoderHeadNet

Multi-layer Perceptron for MNIST. This is a 2-layer neural network for MNIST classification first introduced in the [FedAvg] paper, where the hidden layers have 200 neurons each and the output layer with sigmoid activation.

Similar architectures are also used in other papers:

• [SuPerFed]: hidden_size=(200, 200), same as FedAvg;

• [pFedMe]: hidden_size=(100, 100) with softmax on the output layer;

• [FedDyn]: hidden_size=(200, 100).

Parameters:

• hidden_size (tuple[int, int], optional) – Size of the hidden layers. Defaults to (200, 200).

• softmax (bool, optional) – If True, the output is passed through a softmax layer. Defaults to True.

See also

• MNIST_2NN_E

• MNIST_2NN_D

References

[FedAvg]

H. Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, Blaise Aguera y Arcas. “Communication-Efficient Learning of Deep Networks from Decentralized Data”. In AISTATS (2017).

[SuPerFed]

Seok-Ju Hahn, Minwoo Jeong, and Junghye Lee. Connecting Low-Loss Subspace for Personalized Federated Learning. In KDD (2022).

[pFedMe]

Canh T. Dinh, Nguyen H. Tran, and Tuan Dung Nguyen. Personalized Federated Learning with Moreau Envelopes. In NeurIPS (2020).

[FedDyn]

S. Wang, T. Liu, and M. Hong. “FedDyn: A Dynamic Federated Learning Framework”. In ICLR (2021).

class fluke.nets.MNIST_2NN_E(hidden_size: tuple[int, int] = (200, 100))

Bases: Module

Encoder for the MNIST_2NN network.

Parameters:

hidden_size (tuple[int, int], optional) – Size of the hidden layers. Defaults to (200, 200).

See also

• MNIST_2NN

• MNIST_2NN_D

class fluke.nets.MNIST_2NN_D(hidden_size: int = 100, use_softmax: bool = False)

Bases: Module

Head for the MNIST_2NN network.

Parameters:

• hidden_size (int, optional) – Size of the hidden layer. Defaults to 200.

• use_softmax (bool, optional) – If True, the output is passed through a softmax layer, otherwise, no activation will be used. Defaults to False.

See also

• MNIST_2NN

• MNIST_2NN_E

class fluke.nets.MNIST_CNN

Bases: EncoderHeadNet

Convolutional Neural Network for MNIST. This is a simple CNN for MNIST classification first introduced in the [FedAvg] paper, where the architecture consists of two convolutional layers with 32 and 64 filters, respectively, followed by two fully connected layers with 512 and 10 neurons, respectively.

Very same architecture is also used in the [SuPerFed] paper.

class fluke.nets.MNIST_CNN_E

Bases: Module

Encoder for the MNIST_CNN network.

See also

• MNIST_CNN

• MNIST_CNN_D

class fluke.nets.MNIST_CNN_D

Bases: Module

Head for the MNIST_CNN network.

See also

• MNIST_CNN

• MNIST_CNN_E

class fluke.nets.FedBN_CNN(channels: int = 1)

Bases: EncoderHeadNet

Convolutional Neural Network with Batch Normalization for CIFAR-10. This network follows the architecture proposed in the [FedBN] paper, where the encoder consists of four convolutional layers with 64, 64, 128, and 128 filters, respectively, and the head network consists of three fully connected layers with 2048, 512, and 10 neurons, respectively.

Parameters:

channels (int, optional) – Number of input channels. Defaults to 1.

Note

In the original paper, the size of the last convolutional layer is erroneously reported.

See also

• FedBN_CNN_E

• FedBN_CNN_D

References

[FedBN]

Xiaoxiao Li, Meirui JIANG, Xiaofei Zhang, Michael Kamp, and Qi Dou. FedBN: Federated Learning on Non-IID Features via Local Batch Normalization. In ICLR (2021).

class fluke.nets.FedBN_CNN_E(channels: int = 1)

Bases: Module

Encoder for the FedBN_CNN network.

Parameters:

channels (int, optional) – Number of input channels. Defaults to 1.

See also

• FedBN_CNN

• FedBN_CNN_D

class fluke.nets.FedBN_CNN_D

Bases: Module

Head for the FedBN_CNN network.

See also

• FedBN_CNN

• FedBN_CNN_E

class fluke.nets.MNIST_LR

Bases: Module

Logistic Regression for MNIST. This is a simple logistic regression model for MNIST classification used in the [FedProx] paper for both MNIST and FEMNIST datasets.

References

[FedProx]

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated Optimization in Heterogeneous Networks. Adaptive & Multitask Learning Workshop. In Open Review https://openreview.net/pdf?id=SkgwE5Ss3N (2018).

class fluke.nets.CifarConv2

Bases: EncoderHeadNet

Convolutional Neural Network for CIFAR-10. This is a CNN for CIFAR-10 classification as described in the [FedNH] paper, where the architecture consists of two convolutional layers with 64 filters, followed by two fully connected layers with 384 and 100 neurons, respectively. The convolutional layers are followed by ReLU activations and max pooling. The last classification layer is a linear layer with 10 neurons.

Parameters:

• embedding_size (int, optional) – Size of the embedding after the second linear layer. Defaults to 100.

• num_classes (int, optional) – Number of classes. Defaults to 10.

See also

• CifarConv2_E

• CifarConv2_D

References

[FedNH]

Yutong Dai, Zeyuan Chen, Junnan Li, Shelby Heinecke, Lichao Sun, Ran Xu. Tackling Data Heterogeneity in Federated Learning with Class Prototypes. In AAAI (2023).

class fluke.nets.CifarConv2_E

Bases: Module

Encoder for the CifarConv2 network.

See also

• CifarConv2

• CifarConv2_D

class fluke.nets.CifarConv2_D

Bases: Module

Head for the CifarConv2 network.

See also

• CifarConv2

• CifarConv2_E

class fluke.nets.ResNet9

Bases: EncoderHeadNet

ResNet-9 network for CIFAR-100 classification. This network follows the architecture proposed in the [SuPerFed] paper, which fllows the standard ResNet-9. The encoder consists of all the layers but the last fully connected layer, and thus the head network consists of the last fully connected layer.

See also

• ResNet9_E

• ResNet9_D

class fluke.nets.ResNet9_E

Bases: Module

Encoder for the ResNet9 network.

See also

• ResNet9

• ResNet9_D

class fluke.nets.ResNet9_D

Bases: Module

Head for the ResNet9 network.

See also

• ResNet9

• ResNet9_E

class fluke.nets.FEMNIST_CNN

Bases: EncoderHeadNet

Convolutional Neural Network for FEMNIST. This is a simple CNN for FEMNIST classification first introduced in the [DITTO] paper, where the architecture consists of two convolutional layers with 32 and 64 filters, respectively, followed by two fully connected layers with 1024 and 62 neurons, respectively. Each convolutional layer is followed by a ReLU activation and a max pooling layer.

References

[DITTO]

Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and Robust Federated Learning Through Personalization. In ICML (2021).

class fluke.nets.FEMNIST_CNN_E

Bases: Module

Encoder for the FEMNIST_CNN network.

See also

• FEMNIST_CNN

• FEMNIST_CNN_D

class fluke.nets.FEMNIST_CNN_D

Bases: Module

Head for the FEMNIST_CNN network.

See also

• FEMNIST_CNN

• FEMNIST_CNN_E

class fluke.nets.VGG9_E(seed: int = 98765)

Bases: Module

Encoder for the VGG9 network.

Parameters:

seed (int, optional) – Seed used for weight initialization. Defaults to 98765.

See also

• VGG9

• VGG9_D

class fluke.nets.VGG9_D(input_size: int = 512, output_size: int = 62, seed: int = 98765)

Bases: Module

Head for the VGG9 network.

Parameters:

• input_size (int, optional) – Size of the input tensor. Defaults to 512.

• output_size (int, optional) – Number of output classes. Defaults to 62.

• seed (int, optional) – Seed used for weight initialization. Defaults to 98765.

See also

• VGG9

• VGG9_E

class fluke.nets.VGG9(output_size: int = 62, seed: int = 98765)

Bases: EncoderHeadNet

VGG-9 network for FEMNIST classification. This network follows the architecture proposed in the [SuPerFed] paper which follows the standard VGG-9 architecture. In this implementation all convolutional layers are considered as the encoder and the fully connected layers are considered as the head network.

Parameters:

• input_size (int, optional) – Size of the input tensor. Defaults to 784.

• output_size (int, optional) – Number of output classes. Defaults to 62.

• seed (int, optional) – Seed used for weight initialization. Defaults to 98765.

See also

• VGG9_E

• VGG9_D

class fluke.nets.ResNet18(output_size=10)

Bases: Module

ResNet-18 network as defined in the torchvision library.

Note

This class is a wrapper around the ResNet-18 model from torchvision and it does not implement the EncoderHeadNet interface.

Parameters:

output_size (int, optional) – Number of output classes. Defaults to 10.

class fluke.nets.ResNet34(output_size: int = 100)

Bases: Module

ResNet-34 network as defined in the torchvision library.

Note

This class is a wrapper around the ResNet-18 model from torchvision and it does not implement the EncoderHeadNet interface.

Parameters:

output_size (int, optional) – Number of output classes. Defaults to 100.

class fluke.nets.ResNet50(output_size: int = 100)

Bases: Module

ResNet-50 network as defined in the torchvision library.

Note

This class is a wrapper around the ResNet-18 model from torchvision and it does not implement the EncoderHeadNet interface.

Parameters:

output_size (int, optional) – Number of output classes. Defaults to 100.

class fluke.nets.ResNet18GN(output_size: int = 10)

Bases: ResNet18

ResNet-18 network as defined in the torchvision library but with Group Normalization layers instead of Batch Normalization.

Note

This class is a wrapper around the ResNet-18 model from torchvision and it does not implement the EncoderHeadNet interface.

Parameters:

output_size (int, optional) – Number of output classes. Defaults to 10.

class fluke.nets.MoonCNN_E

Bases: Module

Encoder for the MoonCNN network.

See also

• MoonCNN

• MoonCNN_D

class fluke.nets.MoonCNN_D

Bases: Module

Head for the MoonCNN network.

See also

• MoonCNN

• MoonCNN_E

class fluke.nets.MoonCNN

Bases: EncoderHeadNet

Convolutional Neural Network for CIFAR-10. This is a CNN for CIFAR-10 classification first described in the [MOON] paper, where the architecture consists of two convolutional layers with 6 and 16 filters, respectively, followed by two fully connected layers with 120 and 84 neurons, respectively, and a projection head with 256 neurons followed by the output layer with 10 neurons.

See also

• MoonCNN_E

• MoonCNN_D

References

[MOON]

Qinbin Li, Bingsheng He, and Dawn Song. Model-Contrastive Federated Learning. In CVPR (2021).

class fluke.nets.LeNet5_E

Bases: Module

Encoder for the LeNet5 network.

See also

• LeNet5

• LeNet5_D

class fluke.nets.LeNet5_D(output_size: int = 100)

Bases: Module

Head for the LeNet5 network.

Parameters:

output_size (int, optional) – Number of output classes. Defaults to 100.

See also

• LeNet5

• LeNet5_E

class fluke.nets.LeNet5(output_size: int = 100)

Bases: EncoderHeadNet

LeNet-5 for CIFAR. This is a LeNet-5 for CIFAR-10/100 classification as described in the [FedRep] paper, where the architecture consists of two convolutional layers with 6 and 16 filters, respectively, followed by two fully connected layers with 120 and 84 neurons, respectively, and the output layer with 10 neurons. The activation functions are ReLU. The architecture is also used in the [LG-FedAvg] paper.

See also

• LeNet5_E

• LeNet5_D

Parameters:

output_size (int, optional) – Number of output classes. Defaults to 100.

References

[FedRep]

Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In ICML (2021).

[LG-FedAvg]

Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B. Allen, Randy P. Auerbach, David Brent, Ruslan Salakhutdinov, Louis-Philippe Morency. Think Locally, Act Globally: Federated Learning with Local and Global Representations. In arXiv https://arxiv.org/abs/2001.01523 (2020).

class fluke.nets.Shakespeare_LSTM_E

Bases: Module

Encoder for the Shakespeare_LSTM network.

See also

• Shakespeare_LSTM

• Shakespeare_LSTM_D

class fluke.nets.Shakespeare_LSTM_D

Bases: Module

Head for the Shakespeare_LSTM network.

See also

• Shakespeare_LSTM

• Shakespeare_LSTM_E

class fluke.nets.Shakespeare_LSTM

Bases: EncoderHeadNet

LSTM for Shakespeare. This is an LSTM for Shakespeare classification first introduced in the [SuPerFed] paper, where the architecture consists of an embedding layer with 8 dimensions, followed by a two-layer LSTM with 256 hidden units, and a linear layer with 256 neurons.

See also

• Shakespeare_LSTM_E

• Shakespeare_LSTM_D