Simple Model Architectures

In this file, you can find a number of models that are commonly used in FL community. These models are used in Communication-Efficient Learning of Deep Networks from Decentralized Data.

class SimpleCNN(num_classes=10, num_channels=1, in_height=28, in_width=28, num_filters1=32, num_filters2=64, feature_size=512)

A simple two layer CNN Perceptron.

Parameters

• num_classes (int, optional) -- number of classes. Defaults to 10. Assigning None or a negative integer means no classifier.

• num_channels (int, optional) -- number of channels of input. Defaults to 1.

• in_height (int, optional) -- input height to resize to. Defaults to 28.

• in_width (int, optional) -- input width to resize to. Defaults to 28.

• feature_size (int, optional) -- number of features. Defaults to 512.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_features(x)

Gets the extracted features. Goes through all cells except the classifier.

Parameters

x (Tensor) -- input tensor with shape \((N\times C\times D_1\times D_2\times \dots\times D_n)\) where N is batch size and C is dtermined by num_channels.

Returns

Tensor --

output tensor with shape

\((N\times O)\) where O is determined by feature_size

training: bool

class SimpleCNN2(num_classes=10, num_channels=3, in_height=24, in_width=24, num_filters1=64, num_filters2=64, hidden_size=384, feature_size=192)

A simple two layer CNN Perceptron. This is similar to CNN model in McMahan's FedAvg paper.

Parameters

• num_classes (int, optional) -- number of classes. Defaults to 10. Assigning None or a negative integer means no classifier.

• num_channels (int, optional) -- number of channels of input. Defaults to 1.

• in_height (int, optional) -- input height to resize to. Defaults to 28.

• in_width (int, optional) -- input width to resize to. Defaults to 28.

• hidden_size (int, optional) -- number of hidden neurons. Defaults to 384.

• feature_size (int, optional) -- number of features. Defaults to 192.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_features(x)

Gets the extracted features. Goes through all cells except the classifier.

Parameters

x (Tensor) -- input tensor with shape \((N\times C\times D_1\times D_2\times \dots\times D_n)\) where N is batch size and C is dtermined by num_channels.

Returns

Tensor --

output tensor with shape

\((N\times O)\) where O is determined by feature_size

training: bool

class SimpleMLP(num_classes=10, num_channels=1, in_height=28, in_width=28, feature_size=200)

A simple two layer Multi-Layer Perceptron. This is referred to as 2NN in McMahan's FedAvg paper.

Parameters

• num_classes (int, optional) -- number of classes. Defaults to 10. Assigning None or a negative integer means no classifier.

• num_channels (int, optional) -- number of channels of input. Defaults to 1.

• in_height (int, optional) -- input height to resize to. Defaults to 28.

• in_width (int, optional) -- input width to resize to. Defaults to 28.

• feature_size (int, optional) -- number of features. Defaults to 200.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_features(x)

Gets the extracted features. Goes through all cells except the classifier.

Parameters

x (Tensor) -- input tensor with shape \((N\times C\times D_1\times D_2\times \dots\times D_n)\) where N is batch size and C is dtermined by num_channels.

Returns

Tensor --

output tensor with shape

\((N\times O)\) where O is determined by feature_size

training: bool