-
Notifications
You must be signed in to change notification settings - Fork 1
/
custom_dataset.py
200 lines (167 loc) · 7.18 KB
/
custom_dataset.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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
104
105
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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
"""
The same as FedProxEMNIST, for testing custom dataset.
"""
import json
import warnings
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union
import numpy as np
import torch
import torch.utils.data as torchdata
from fl_sim.data_processing._register import register_fed_dataset
from fl_sim.data_processing.fed_dataset import FedVisionDataset
from fl_sim.models import nn as mnn
from fl_sim.models.utils import top_n_accuracy
from fl_sim.utils._download_data import url_is_reachable
from fl_sim.utils.const import CACHED_DATA_DIR
__all__ = [
"CustomFEMNIST",
]
CUSTOM_FEMNIST_DATA_DIR = CACHED_DATA_DIR / "custom_femnist"
CUSTOM_FEMNIST_DATA_DIR.mkdir(parents=True, exist_ok=True)
_label_mapping = {i: c for i, c in enumerate("abcdefghijklmnopqrstuvwxyz"[:10])}
@register_fed_dataset()
class CustomFEMNIST(FedVisionDataset):
"""
(part of) federeated EMNIST used in the FedProx paper and the FedDR paper.
`The FEMNIST data we used in the paper is a subsampled (and repartitioned) version
of the original full dataset in order to impose additional statistical heterogeneity.
The above dataset is generated by the following instruction: ...`
NOTE that this dataset is not the same as the original FEMNIST dataset,
containing only 10 classes (a-j), instead of 62 classes (a-z, A-Z, 0-9).
The raw data has been processed using min-max normalization to range [0, 1],
hence any further augmentation is perhaps inappropriate.
Parameters
----------
datadir : Union[Path, str], optional
Directory to store data.
If ``None``, use default directory.
transform : Union[str, Callable], default "none"
Transform to apply to data. Conventions:
``"none"`` means no transform, using TensorDataset.
seed : int, default 0
Random seed for data partitioning.
**extra_config : dict, optional
Extra configurations.
References
----------
.. [1] https://github.com/litian96/FedProx/tree/master/data/nist
"""
__name__ = "CustomFEMNIST"
def _preload(self, datadir: Optional[Union[str, Path]] = None) -> None:
self.datadir = Path(datadir or CUSTOM_FEMNIST_DATA_DIR).expanduser().resolve()
self.DEFAULT_TRAIN_CLIENTS_NUM = 200
self.DEFAULT_TEST_CLIENTS_NUM = 200
self.DEFAULT_BATCH_SIZE = 20
self.DEFAULT_TRAIN_FILE = Path("train") / Path("mytrain.json")
self.DEFAULT_TEST_FILE = Path("test") / Path("mytest.json")
self._EXAMPLE = "user_data"
self._IMGAE = "x"
self._LABEL = "y"
if self.transform != "none":
warnings.warn(
"The images are not raw pixels, but processed. " "The transform argument will be ignored.",
RuntimeWarning,
)
self.transform = "none"
self.criterion = torch.nn.CrossEntropyLoss()
self.download_if_needed()
# client id list
train_file_path = self.datadir / self.DEFAULT_TRAIN_FILE
test_file_path = self.datadir / self.DEFAULT_TEST_FILE
self._train_data_dict = json.loads(train_file_path.read_text())
self._test_data_dict = json.loads(test_file_path.read_text())
self._client_ids_train = self._train_data_dict["users"]
self._client_ids_test = self._test_data_dict["users"]
self._n_class = len(
np.unique(
np.concatenate(
[
self._train_data_dict[self._EXAMPLE][self._client_ids_train[idx]][self._LABEL]
for idx in range(self.DEFAULT_TRAIN_CLIENTS_NUM)
]
)
)
)
def get_dataloader(
self,
train_bs: Optional[int] = None,
test_bs: Optional[int] = None,
client_idx: Optional[int] = None,
) -> Tuple[torchdata.DataLoader, torchdata.DataLoader]:
# load data
if client_idx is None:
# get ids of all clients
train_ids = self._client_ids_train
test_ids = self._client_ids_test
else:
# get ids of single client
train_ids = [self._client_ids_train[client_idx]]
test_ids = [self._client_ids_test[client_idx]]
# load data
train_x = np.vstack([self._train_data_dict[self._EXAMPLE][client_id][self._IMGAE] for client_id in train_ids])
train_y = np.concatenate([self._train_data_dict[self._EXAMPLE][client_id][self._LABEL] for client_id in train_ids])
test_x = np.vstack([self._test_data_dict[self._EXAMPLE][client_id][self._IMGAE] for client_id in test_ids])
test_y = np.concatenate([self._test_data_dict[self._EXAMPLE][client_id][self._LABEL] for client_id in test_ids])
# dataloader
train_ds = torchdata.TensorDataset(
torch.from_numpy(train_x.reshape((-1, 28, 28)).astype(np.float32)).unsqueeze(1),
torch.from_numpy(train_y.astype(np.int64)),
)
train_dl = torchdata.DataLoader(
dataset=train_ds,
batch_size=train_bs or self.DEFAULT_BATCH_SIZE,
shuffle=True,
drop_last=False,
)
test_ds = torchdata.TensorDataset(
torch.from_numpy(test_x.reshape((-1, 28, 28)).astype(np.float32)).unsqueeze(1),
torch.from_numpy(test_y.astype(np.int64)),
)
test_dl = torchdata.DataLoader(
dataset=test_ds,
batch_size=test_bs or self.DEFAULT_BATCH_SIZE,
shuffle=True,
drop_last=False,
)
return train_dl, test_dl
def extra_repr_keys(self) -> List[str]:
return [
"n_class",
] + super().extra_repr_keys()
def evaluate(self, probs: torch.Tensor, truths: torch.Tensor) -> Dict[str, float]:
return {
"acc": top_n_accuracy(probs, truths, 1),
"top3_acc": top_n_accuracy(probs, truths, 3),
"top5_acc": top_n_accuracy(probs, truths, 5),
"loss": self.criterion(probs, truths).item(),
"num_samples": probs.shape[0],
}
@property
def url(self) -> str:
# https://drive.google.com/file/d/1tCEcJgRJ8NdRo11UJZR6WSKMNdmox4GC/view?usp=sharing
# "http://218.245.5.12/NLP/federated/fedprox-femnist.zip"
if url_is_reachable("http://www.dropbox.com"):
return "https://www.dropbox.com/s/55ibep82qqars9w/fedprox-femnist.zip?dl=1"
else:
return "https://deep-psp.tech/Data/FL/fedprox-femnist.zip"
@property
def candidate_models(self) -> Dict[str, torch.nn.Module]:
"""A set of candidate models."""
return {
"cnn_femmist_tiny": mnn.CNNFEMnist_Tiny(num_classes=self.n_class),
"cnn_femmist": mnn.CNNFEMnist(num_classes=self.n_class),
# "resnet10": mnn.ResNet10(num_classes=self.n_class),
"mlp": mnn.MLP(dim_in=28 * 28, dim_out=self.n_class, ndim=2),
}
@property
def doi(self) -> List[str]:
return [
"10.1109/5.726791", # MNIST
"10.1109/ijcnn.2017.7966217", # EMNIST
"10.48550/ARXIV.1812.01097", # LEAF
"10.48550/ARXIV.1812.06127", # FedProx
]
@property
def label_map(self) -> dict:
return _label_mapping