Skip to content

taps.apps.fedlearn.modules

taps/apps/fedlearn/modules.py

CifarModule 

CifarModule(num_classes: int)

Bases: Module

Cifar model.

Source: https://www.kaggle.com/code/shadabhussain/cifar-10-cnn-using-pytorch

Source code in taps/apps/fedlearn/modules.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(self, num_classes: int):
    super().__init__()
    self.num_classes = num_classes
    self.network = nn.Sequential(
        nn.Conv2d(3, 32, kernel_size=3, padding=1),
        nn.ReLU(),
        nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(2, 2),  # output: 64 x 16 x 16
        nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(2, 2),  # output: 128 x 8 x 8
        nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
        nn.ReLU(),
        nn.MaxPool2d(2, 2),  # output: 256 x 4 x 4
        nn.Flatten(),
        nn.Linear(256 * 4 * 4, 1024),
        nn.ReLU(),
        nn.Linear(1024, 512),
        nn.ReLU(),
        nn.Linear(512, num_classes),
    )

forward() 

forward(x: Tensor) -> Tensor

Forward pass.

Source code in taps/apps/fedlearn/modules.py 
49
50
51
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass."""
    return self.network(x)

MnistModule 

MnistModule()

Bases: Module

Model for MNIST and FashionMNIST data.

Source code in taps/apps/fedlearn/modules.py 
57
58
59
60
61
62
63
def __init__(self) -> None:
    super().__init__()
    self.flattener = nn.Flatten()
    self.fc1 = nn.Linear(28 * 28, 56 * 56)
    self.fc2 = nn.Linear(56 * 56, 28 * 28)
    self.fc3 = nn.Linear(28 * 28, 14 * 14)
    self.classifier = nn.Linear(14 * 14, 10)

forward() 

forward(x: Tensor) -> Tensor

Forward pass.

Source code in taps/apps/fedlearn/modules.py 
65
66
67
68
69
70
71
72
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass."""
    x = self.flattener(x)
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = F.relu(self.fc3(x))
    x = self.classifier(x)
    return x

create_model() 

create_model(data: DataChoices) -> Module

Create a model suitable for the dataset choice.

Note

The currently supported dataset options are MNIST, FashionMNIST, CIFAR10, and CIFAR100.

Parameters:

  • data (DataChoices) –

    Name of dataset that will be used for training (and testing).

Returns:

  • Module –

    PyTorch model.

Raises:

  • ValueError –

    If an unsupported value for data is provided.

Source code in taps/apps/fedlearn/modules.py 
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def create_model(data: DataChoices) -> nn.Module:
    """Create a model suitable for the dataset choice.

    Note:
        The currently supported dataset options are `MNIST`, `FashionMNIST`,
        `CIFAR10`, and `CIFAR100`.

    Args:
        data: Name of dataset that will be used for training (and testing).

    Returns:
        PyTorch model.

    Raises:
        ValueError: If an unsupported value for `data` is provided.
    """
    name = data.value.lower()

    if name == 'cifar10':
        return CifarModule(10)
    elif name == 'cifar100':
        return CifarModule(100)
    elif name in ('fmnist', 'mnist'):
        return MnistModule()
    else:
        raise ValueError(
            f'Unknown dataset "{data.value}". Supported options are '
            "'cifar10', 'cifar100', 'fmnist', and 'mnist'.",
        )

load_data() 

load_data(
    data_name: DataChoices,
    root: Path,
    train: bool,
    download: bool = False,
) -> Dataset

Load dataset for training.

Parameters:

  • data_name (DataChoices) –

    Dataset choice.

  • root (Path) –

    Root dataset directory.

  • train (bool) –

    Flag for if training.

  • download (bool, default: False ) –

    Should the dataset be downloaded.

Returns:

  • Dataset ( Dataset ) –

    description

Source code in taps/apps/fedlearn/modules.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def load_data(
    data_name: DataChoices,
    root: pathlib.Path,
    train: bool,
    download: bool = False,
) -> Dataset:
    """Load dataset for training.

    Args:
        data_name: Dataset choice.
        root: Root dataset directory.
        train: Flag for if training.
        download: Should the dataset be downloaded.

    Returns:
        Dataset: _description_
    """
    kwargs = {
        'root': root,
        'train': train,
        'transform': transforms.ToTensor(),
        'download': download,
    }
    name = data_name.value.lower()
    if name == 'cifar10':
        return torchvision.datasets.CIFAR10(**kwargs)
    elif name == 'cifar100':
        return torchvision.datasets.CIFAR100(**kwargs)
    elif name == 'fmnist':
        return torchvision.datasets.FashionMNIST(**kwargs)
    elif name == 'mnist':
        return torchvision.datasets.MNIST(**kwargs)
    else:
        raise ValueError(f'Unknown dataset: {data_name}.')