p2pfl.learning.frameworks.pytorch.lightning_model module

Convolutional Neural Network (for MNIST) with PyTorch Lightning.

class p2pfl.learning.frameworks.pytorch.lightning_model.LightningModel(model, params=None, num_samples=None, contributors=None, additional_info=None)

Bases: P2PFLModel

P2PFL model abstraction for PyTorch Lightning.

Parameters:

• model (LightningModule) – The model to encapsulate.

• params (Union[List[ndarray], bytes, None]) – The parameters of the model.

• num_samples (Optional[int]) – The number of samples.

• contributors (Optional[List[str]]) – The contributors of the model.

• additional_info (Optional[Dict[str, Any]]) – Additional information.

get_framework()

Retrieve the model framework name.

Return type:

str

Returns:

The name of the model framework.

get_parameters()

Get the parameters of the model.

Return type:

List[ndarray]

Returns:

The parameters of the model

set_parameters(params)

Set the parameters of the model.

Parameters:

params (Union[List[ndarray], bytes]) – The parameters of the model.

Raises:

ModelNotMatchingError – If parameters don’t match the model.

Return type:

None

class p2pfl.learning.frameworks.pytorch.lightning_model.MLP(input_size=784, hidden_sizes=None, out_channels=10, activation='relu', metric=<class 'torchmetrics.classification.accuracy.Accuracy'>, lr_rate=0.001, seed=None)

Bases: LightningModule

Multilayer Perceptron (MLP) with configurable parameters.

configure_optimizers()

Configure the optimizer.

Return type:

Optimizer

forward(x)

Forward pass of the MLP.

Return type:

Tensor

test_step(batch, batch_id)

Test step for the MLP.

Return type:

Tensor

training_step(batch, batch_id)

Training step of the MLP.

Return type:

Tensor

validation_step(batch, batch_id)

Perform validation step for the MLP.

Return type:

Tensor