fedbox.datasets.cifar10

 1from torchvision.datasets import CIFAR10
 2from typing import Any, Union
 3
 4from . import utils
 5
 6
 7def cifar10(
 8    directory: str,
 9    n_subsets: int = 100,
10    heterogeneity_degree: float = None,
11    samples_variance: float = 0.0,
12    transform: Any = None,
13    target_transform: Any = None,
14    seed: int = None,
15    strict: bool = True
16) -> dict[str, list[utils.FederatedSubset]]:
17    '''
18    Loads the `CIFAR10` dataset and partitions it into `n_subsets` training subsets and 
19    `n_subsets` testing subsets according to heterogeneity parameters `heterogeneity_degree` 
20    and `samples_variance`.
21
22    Parameters
23    ----------
24    directory: str
25        Root directory to load the dataset
26    n_subsets: int
27        Number of datasets for splitting
28    heterogeneity_degree: float
29        Class heterogeneity degree, by default is homogeneous
30    samples_variance: float
31        Standard deviation (%) in the number of samples for each client, `0` by default
32    transform: Any
33        Transformation to apply on data samples, `None` by default
34    target_transform: Any
35       Transformation to apply on data labels, `None` by default
36    seed: int
37        Random seed initializer
38    strict: bool
39        In strict mode `heterogeneity_degree` and `samples_variance` are highly respected,
40        otherwise flexibility is allowed
41
42    Returns
43    -------
44    dict[str, list[utils.FederatedSubset]]
45        Returns the lists of subsets of training clients and testing clients
46    '''
47
48    training = CIFAR10(
49        root = directory, 
50        train = True, 
51        download = True,
52        transform = transform,
53        target_transform = target_transform
54    )
55
56    testing = CIFAR10(
57        root = directory, 
58        train = False, 
59        download = True,
60        transform = transform,
61        target_transform = target_transform
62    )
63
64    return {
65        'training': utils.partition(
66            training,
67            n_subsets = min(len(training), n_subsets),
68            n_classes = 10,
69            heterogeneity_degree = heterogeneity_degree,
70            samples_variance = samples_variance,
71            return_indices = False,
72            seed = seed,
73            strict = strict
74        ),
75        'testing': utils.partition(
76            testing,
77            n_subsets = min(len(testing), n_subsets),
78            n_classes = 10,
79            heterogeneity_degree = heterogeneity_degree,
80            samples_variance = samples_variance,
81            return_indices = False,
82            seed = seed,
83            strict = strict
84        )
85    }