p2pfl.learning.frameworks.pytorch.lightning_dataset module

PyTorch dataset integration.

class p2pfl.learning.frameworks.pytorch.lightning_dataset.PyTorchExportStrategy[source]

Bases: DataExportStrategy

Export strategy for PyTorch tensors.

static export(data, batch_size=None, num_workers=0, **kwargs)[source]

Export the data using the PyTorch strategy.

Parameters:

  • data (Dataset) – The data to export. Transforms should already be applied to the dataset via set_transform.

  • batch_size (int | None) – The batch size to use for the exported data.

  • num_workers (int) – The number of workers to use for the exported

  • kwargs – Additional keyword arguments.

Return type:

DataLoader

Returns:

The exported data.

class p2pfl.learning.frameworks.pytorch.lightning_dataset.TorchvisionDatasetFactory[source]

Bases: object

Factory class for loading PyTorch Vision datasets in P2PFL.

static get_mnist(cache_dir, train=True, download=True)[source]

Get the MNIST dataset from PytorchVision.

Parameters:

  • cache_dir (str | Path) – The directory where the dataset will be stored.

  • train (bool) – Whether to get the training or test dataset.

  • download (bool) – Whether to download the dataset.

Return type:

P2PFLDataset