Editing data labels

Hello,

I am trying to split Cifar10 to get two subsets (Animals and Vehicles). I was able to successfully separate them by the labels

partitioner = IidPartitioner(num_partitions=1)
fds = FederatedDataset(dataset="uoft-cs/cifar10", partitioners={"train": partitioner},)

partition1_train_test = fds.load_partition(0).train_test_split(test_size=0.2, seed=42)

pytorch_transforms = Compose([ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

def apply_transforms(batch):
    """Apply transforms to the partition from FederatedDataset."""
    batch["img"] = [pytorch_transforms(img) for img in batch["img"]]
    return batch

partition_train_test = partition_train_test.with_transform(apply_transforms)

labels_list_partition1_train = [item['label'] for item in partition_train_test['train']]

animals = torch.tensor([2, 3, 4, 5, 6, 7])
animalsindicesTrain = (torch.tensor(labels_list_partition1_train)[..., None] == animals).any(-1).nonzero(as_tuple=True)[0]

animalsTrainSubset = torch.utils.data.Subset(partition1_train_test['train'], animalsindicesTrain)

animalsTrainDataloader = DataLoader(animalsTrainSubset, batch_size=128, shuffle=True, num_workers=2)

But I still have 10 classes in that subset and I would like it to have 6 for animals and 4 for vehicles. Is in my interest to have different length heads as I am doing some tests with federated environments.

Is there any straight forward way of reducing the number of classes of the subsets?

Thank you in advance.

Subset is only a thin wrapper, which uses the specified indices and forwards it to the underlying Dataset. The same Dataset.__getitem__ (with a subset of indices) will be used to load and transform the samples.

You could try to remap the target indices to new ranges (both starting at 0 to avoid issues with your criterion) e.g. via:

class RemappedLabels(torch.utils.data.Dataset):
    def __init__(self, dataset, label_map):
        self.dataset = dataset
        self.label_map = label_map
    def __len__(self):
        return len(self.dataset)
    def __getitem__(self, idx):
        sample = dict(self.dataset[idx])
        sample["label"] = self.label_map[int(sample["label"])]
        return sample
...
animal_map = {2: 0, 3: 1, 4: 2, 5: 3, 6: 4, 7: 5}
vehicle_map = {0: 0, 1: 1, 8: 2, 9: 3}
animalsTrainSubset = torch.utils.data.Subset(partition1_train_test["train"], animalsindicesTrain)
animalsTrainDataset = RemappedLabels(animalsTrainSubset, animal_map)

Let me know if this would work for your use case.

Quick and straightforward, that worked at first try. Thank you very much! Now I can continue with my investigations.