Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix divide_dataset in Federated Datasets #3192

Merged
merged 4 commits into from
Apr 3, 2024
Merged

Fix divide_dataset in Federated Datasets #3192

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ab6aef8
Add test for more than 2 divisions
adam-narozniak Apr 2, 2024
9c7fc6b
Fix more than 2 divisions
adam-narozniak Apr 2, 2024
c768f3d
Fix formatting
adam-narozniak Apr 2, 2024
380b7d3
Merge branch 'main' into fds-fix-divide-dataset
jafermarq Apr 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions datasets/flwr_datasets/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,7 @@ def divide_dataset(
>>> train_test = divide_dataset(dataset=partition, division=division)
>>> train, test = train_test["train"], train_test["test"]
"""
_check_division_config_correctness(division)
dataset_length = len(dataset)
ranges = _create_division_indices_ranges(dataset_length, division)
if isinstance(division, (list, tuple)):
Expand Down Expand Up @@ -162,15 +163,15 @@ def _create_division_indices_ranges(
for fraction in division:
end_idx += int(dataset_length * fraction)
ranges.append(range(start_idx, end_idx))
start_idx += end_idx
start_idx = end_idx
elif isinstance(division, dict):
ranges = []
start_idx = 0
end_idx = 0
for fraction in division.values():
end_idx += int(dataset_length * fraction)
ranges.append(range(start_idx, end_idx))
start_idx += end_idx
start_idx = end_idx
else:
TypeError(
f"The type of the `division` should be dict, "
Expand Down Expand Up @@ -274,6 +275,7 @@ def concatenate_divisions(
concatenated_divisions : Dataset
A dataset created as concatenation of the divisions from all partitions.
"""
_check_division_config_correctness(partition_division)
divisions = []
zero_len_divisions = 0
for partition_id in range(partitioner.num_partitions):
Expand Down
23 changes: 21 additions & 2 deletions datasets/flwr_datasets/utils_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,13 +31,32 @@
"expected_concatenation_size",
),
[
# Create 1 division
((1.0,), [40], 0, 40),
({"train": 1.0}, [40], "train", 40),
# Create 2 divisions
((0.8, 0.2), [32, 8], 1, 8),
([0.8, 0.2], [32, 8], 1, 8),
({"train": 0.8, "test": 0.2}, [32, 8], "test", 8),
# Create 3 divisions
([0.6, 0.2, 0.2], [24, 8, 8], 1, 8),
({"train": 0.6, "valid": 0.2, "test": 0.2}, [24, 8, 8], "test", 8),
# Create 4 divisions
([0.4, 0.2, 0.2, 0.2], [16, 8, 8, 8], 1, 8),
({"0": 0.4, "1": 0.2, "2": 0.2, "3": 0.2}, [16, 8, 8, 8], "1", 8),
# Not full dataset
# Create 1 division
([0.8], [32], 0, 32),
({"train": 0.8}, [32], "train", 32),
# Create 2 divisions
([0.2, 0.1], [8, 4], 1, 4),
((0.2, 0.1), [8, 4], 0, 8),
({"train": 0.2, "test": 0.1}, [8, 4], "test", 4),
# Create 3 divisions
([0.6, 0.2, 0.1], [24, 8, 4], 2, 4),
({"train": 0.6, "valid": 0.2, "test": 0.1}, [24, 8, 4], "test", 4),
# Create 4 divisions
([0.4, 0.2, 0.1, 0.2], [16, 8, 4, 8], 2, 4),
({"0": 0.4, "1": 0.2, "2": 0.1, "3": 0.2}, [16, 8, 4, 8], "2", 4),
],
)
class UtilsTests(unittest.TestCase):
Expand All @@ -60,7 +79,7 @@ def test_correct_sizes(self) -> None:
else:
lengths = [len(split) for split in divided_dataset.values()]

self.assertEqual(lengths, self.sizes)
self.assertEqual(self.sizes, lengths)

def test_correct_return_types(self) -> None:
"""Test correct types of the divided dataset based on the config."""
Expand Down