dataset.py

from operator import le
import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets

from sklearn.model_selection import train_test_split
from sklearn.model_selection import StratifiedShuffleSplit

from resnet import custom_cnn


def get_cifar10_iid_train_loader(data_dir, batch_size, random_seed, shuffle=True, num_workers=4, pin_memory=True):
    # define transforms
    trans = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(degrees=15),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # load dataset
    dataset = datasets.CIFAR10(root=data_dir, transform=trans, download=True, train=True)

    if shuffle:
        np.random.seed(random_seed)

    is_iid, pnumber = True, 5
    coefs = [0.8, 0.75, 2 / 3, 0.5]
    lst = []
    temp = dataset
    while len(lst) < pnumber:
        labels = [i[1] for i in temp]
        t1_set, t2_set = train_test_split(
            temp,
            test_size=coefs[len(lst)],
            random_state=101,
            stratify=labels
        )
        lst.append(t1_set)
        if len(lst) == pnumber - 1:
            lst.append(t2_set)
        temp = t2_set

    t_loaders = [torch.utils.data.DataLoader(elem, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers,
                                             pin_memory=pin_memory) for elem in lst]

    train_loader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory,
    )
    print(t_loaders)
    return t_loaders + [train_loader]


def get_cifar10_noniid_train_loader(data_dir, batch_size, random_seed, shuffle=True, num_workers=4, pin_memory=True):
    # define transforms
    trans = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(degrees=15),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # load dataset
    dataset = datasets.CIFAR10(root=data_dir, transform=trans, download=True, train=True)

    if shuffle:
        np.random.seed(random_seed)
    torch.manual_seed(101)

    is_iid, pnumber = True, 5
    coefs = [0.8, 0.75, 2 / 3, 0.5]
    # coefs = [0.9, 8 / 9, 7 / 8, 6 / 7, 5 / 6, 0.8, 0.75, 2 / 3, 0.5]

    lst = []
    temp = dataset
    while len(lst) < pnumber:
        labels = [i[1] for i in temp]
        t1_set, t2_set = train_test_split(
            temp,
            test_size=coefs[len(lst)],
            random_state=101,
            stratify=labels
        )
        lst.append(t1_set)
        if len(lst) == pnumber - 1:
            lst.append(t2_set)
        temp = t2_set

    print('length of t1 set:', len(lst[0]))
    print('length of t2 set:', len(lst[1]))

    counts = [0] * 10
    for label in lst[0]:
        counts[label[1]] += 1
    print('first set: ', counts, end="\n\n")

    t_loaders = [torch.utils.data.DataLoader(elem, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers,
                                             pin_memory=pin_memory) for elem in lst]

    train_loader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory,
    )
    print(t_loaders)
    return t_loaders + [train_loader]


def get_mnist_iid_train_loader(data_dir, batch_size, random_seed, shuffle=True, num_workers=4, pin_memory=True):
    # define transforms
    trans = transforms.Compose([
        transforms.ToTensor(),
    ])

    # load dataset
    dataset = datasets.MNIST(root=data_dir,
                             transform=trans,
                             download=True,
                             train=True)
    if shuffle:
        np.random.seed(random_seed)

    is_iid, pnumber = True, 2
    if is_iid:
        coefs = [0.8, 0.75, 2 / 3, 0.5] if pnumber == 5 else [0.5]
        lst = []
        temp = dataset 
        while len(lst) < pnumber:
            labels = [i[1] for i in temp]
            t1_set, t2_set = train_test_split(
                temp,
                test_size=coefs[len(lst)],
                random_state=101,
                stratify=labels
            )
            lst.append(t1_set)
            if len(lst) == pnumber - 1:
                lst.append(t2_set)
            temp = t2_set

        counts = [0] * 10
        for label in lst[0]:
            counts[label[1]] += 1
        print('first set: ', counts, end="\n\n")

        t_loaders = [torch.utils.data.DataLoader(elem, batch_size=batch_size, shuffle=shuffle, 
                        num_workers=num_workers, pin_memory=pin_memory) for elem in lst]


    train_loader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=pin_memory,
    )

    return t_loaders + [train_loader]


def main():
    batch_size = 128
    num_workers = 4
    shuffle = True
    pin_memory = True
    if shuffle:
        np.random.seed(101)

    # define transforms
    train_trans = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(degrees=15),
        transforms.ToTensor(),
        transforms.Normalize([0.4915, 0.4823, 0.4468], [0.2470, 0.2435, 0.2616])
    ])

    test_trans = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.4915, 0.4823, 0.4468], [0.2470, 0.2435, 0.2616])
    ])

    # load dataset 
    trainset = torchvision.datasets.CIFAR10(root='data/cifar10/', train=True, download=True, transform=train_trans)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers,
                                              pin_memory=pin_memory)

    testset = torchvision.datasets.CIFAR10(root='data/cifar10/', train=False, download=True, transform=test_trans)
    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=num_workers,
                                             pin_memory=pin_memory)

    print('train:', trainset)
    print('test:', testset, end="\n\n")

    print('--------------------------------------------------------')

    is_iid, pnumber = True, 2
    dataset_size = 25000
    num_classes = len(set(trainset.targets))
    class_size = dataset_size // num_classes # 2500 to pick for each participant
    
    if is_iid:
        # labels
        labels = trainset.targets
        dct = {}
        for idx, label in enumerate(labels):
            if label not in dct:
                dct[label] = []
            dct[label].append(idx)
        
        # division
        lst = {i: [] for i in range(pnumber)}
        for pid in range(pnumber):
            for idx, label in enumerate(dct):
                lst[pid] += list(np.random.choice(dct[idx], size=class_size))
        print(lst[0][:10])
        print(len(lst[0]))
        print(lst[1][:10])


        subsets = [torch.utils.data.Subset(trainset, lst[i]) for i in range(pnumber)]
        t_loaders = [torch.utils.data.DataLoader(subsets[i], shuffle=True, batch_size=batch_size) for i in range(pnumber)]

        counts = [0] * 10
        for label in subsets[0]:
            counts[label[1]] += 1
        print('first set: ', counts, sum(counts), end="\n\n")




    # shuffled_indices = torch.randperm(len(trainset))
    # training_inputs = trainset.data[shuffled_indices]
    # training_labels = torch.Tensor(trainset.targets)[shuffled_indices]
    # split_size = len(trainset) // 5

    # lst = list(
    #     zip(
    #         torch.split(torch.Tensor(training_inputs), split_size),
    #         torch.split(torch.Tensor(training_labels), split_size)
    #     )
    # )

    # # print(lst)
    # print(len(lst))
    # print(lst[0][0].shape)
    # print(lst[0][1].shape)

    # t_loaders = [torch.utils.data.DataLoader(elem, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers,
    #                                          pin_memory=pin_memory) for elem in lst]

    # print(t_loaders[0].dataset)
    # print(t_loaders[0].dataset[1])
    # # print(t_loaders)

    # # finalize bunches of local datasets
    # print(len(lst[0][1]))

    # images, labels = [i[0] for i in trainset], [i[1] for i in trainset]

    # t1_set, t2_set = train_test_split(
    #     trainset,
    #     test_size=0.5,
    #     train_size=0.5,
    #     random_state=101,
    #     stratify=labels
    # )

    # print('length of t1 set:', len(t1_set))
    # print('length of t2 set:', len(t2_set))

    # counts = [0] * 10
    # for label in t1_set:
    #     counts[label[1]] += 1
    # print('first set: ', counts, end="\n\n")

    # counts = [0] * 10
    # for label in t2_set:
    #     counts[label[1]] += 1
    # print('second set: ', counts, end="\n\n")

    # t1_loader = torch.utils.data.DataLoader(t1_set, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers,
    #                                         pin_memory=pin_memory)
    # t2_loader = torch.utils.data.DataLoader(t2_set, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers,
    #                                         pin_memory=pin_memory)

    # net = custom_cnn().cuda()
    # criterion = nn.CrossEntropyLoss().cuda()
    # # optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
    # optimizer = optim.Adam(net.parameters(), lr=0.001)

    # print('net:', net, end="\n\n")

    # for epoch in range(5):
    #     print(epoch)
    #     running_loss = 0.0
    #     for i, data in enumerate(t1_loader, 0):
    #         inputs, labels = data

    #         optimizer.zero_grad()

    #         outputs = net(inputs.cuda())
    #         loss = criterion(outputs, labels.cuda())
    #         loss.backward()
    #         optimizer.step()

    #         # print statistics
    #         running_loss += loss.item()
    #         if i % 20 == 19:  # print every 50 mini-batches
    #             print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
    #             running_loss = 0.0

    # correct = 0
    # total = 0
    # # since we're not training, we don't need to calculate the gradients for our outputs
    # with torch.no_grad():
    #     for data in testloader:
    #         images, labels = data
    #         # calculate outputs by running images through the network
    #         outputs = net(images.cuda())
    #         # the class with the highest energy is what we choose as prediction
    #         _, predicted = torch.max(outputs.data.detach().cpu(), 1)
    #         total += labels.size(0)
    #         correct += (predicted == labels).sum().item()

    # print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')


if __name__ == '__main__':
    main()