fedsim.data_manager.utils — fedsim 0.1 documentation

# Source code for fedsim.data_manager.utils

    import numpy as np
    from torch.utils import data

    class Subset(data.Dataset):
        r"""Subset of a dataset at specified indices.

        Args:
            dataset (Dataset): The whole Dataset
            indices (sequence): Indices in the whole set selected for subset.
        """

        def __init__(self, dataset, indices, transform=None):
            self.dataset = dataset
            if isinstance(indices, int) and indices == -1:
                self.indices = range(len(dataset))
            else:
                self.indices = indices
            self.transform = transform
            targets = np.array(dataset.targets)
            self.targets = targets[self.indices]
            # remove the transform function of the original dataset if transform
            # is provided avoiding double transform
            if transform is not None and self.dataset.transform is not None:
                self.dataset.transform = None

        def __getitem__(self, idx):
            if isinstance(idx, list):
                x, y = self.dataset[[self.indices[i] for i in idx]]
                if self.transform is None:
                    return x, y
                return self.transform(x), y
            x, y = self.dataset[self.indices[idx]]
            if self.transform is None:
                return x, y
            return self.transform(x), y

        def __len__(self):
            return len(self.indices)