Exemplo n.º 1
0
def partition_dataset(args):
    if args.dataset == 'mnist':
        trans_mnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
        dataset_train = datasets.MNIST('../data/mnist/', train=True, download=True, transform=trans_mnist)
        dataset_test = datasets.MNIST('../data/mnist/', train=False, download=True, transform=trans_mnist)
        # sample users
        if args.iid:
            dict_users = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users = mnist_noniid(dataset_train, args.num_users)
    elif args.dataset == 'cifar':
        trans_cifar = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        dataset_train = datasets.CIFAR10('../data/cifar', train=True, download=True, transform=trans_cifar)
        dataset_test = datasets.CIFAR10('../data/cifar', train=False, download=True, transform=trans_cifar)
        if args.iid:
            dict_users = cifar_iid(dataset_train, args.num_users)
        else:
            exit('Error: only consider IID setting in CIFAR10')
    elif args.dataset == 'traffic':
        dataset_train, dataset_test = get_train_valid_loader(
            '/Users/austineyapp/Documents/REP/Year_4/FYP/FederatedLearning/Federated-Learning-for-Traffic-Sign-Recognition',
            batch_size=32, num_workers=0)
        if args.iid:
            dict_users = traffic_iid(dataset_train, args.num_users)
        else:
            exit('Error: only consider IID setting in Traffic')
    else:
        exit('Error: unrecognized dataset')
    return dataset_train, dataset_test, dict_users
Exemplo n.º 2
0
def gen_data(iid, dataset_type, data_root_dir, transforms_dict, num_users, cluster, dataset_split=''):
    # load dataset 
    _, image_datasets, dataset_sizes, class_names =\
            load_dataset(dataset_type, data_root_dir, transforms_dict, batch_size=8, shuffle_flag=False, dataset_split=dataset_split)
    
    dataset_train = image_datasets['train']
    dataset_test = image_datasets['test']
    
    if dataset_type in ['mnist', 'MNIST']:
        # sample users
        if iid:
            dict_users = mnist_iid(dataset_train, num_users)
        else:
            dict_users = mnist_noniid_cluster(dataset_train, num_users, cluster)
    #
    elif dataset_type in ['emnist', 'EMNIST']:     
        if not iid:
            dict_users = emnist_noniid_cluster(dataset_train, num_users, cluster, 
                                               random_shuffle=True)
    #       
    elif dataset_type in ['cifar', 'CIFAR10']:
        if iid:
            dict_users = cifar_iid(dataset_train, num_users)
        else:
            dict_users = cifar_noniid_cluster(dataset_train, num_users, cluster)
    #
    else:
        exit('Error: unrecognized dataset')

    return dataset_train, dataset_test, dict_users
Exemplo n.º 3
0
def partition_dataset():
    """
    :return
    dataset_train: dataset
    dict_users: list of data index for each user. e.g. 100 indexes with 600 data points index
    """
    # load dataset and split users
    if args.dataset == 'mnist':
        trans_mnist = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307, ), (0.3081, ))
        ])
        dataset_train = datasets.MNIST('../data/mnist/',
                                       train=True,
                                       download=True,
                                       transform=trans_mnist)
        dataset_test = datasets.MNIST('../data/mnist/',
                                      train=False,
                                      download=True,
                                      transform=trans_mnist)
        # sample users
        if args.iid:
            dict_users = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users = mnist_noniid(dataset_train, args.num_users)
    elif args.dataset == 'cifar':
        trans_cifar = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        dataset_train = datasets.CIFAR10('../data/cifar',
                                         train=True,
                                         download=True,
                                         transform=trans_cifar)
        dataset_test = datasets.CIFAR10('../data/cifar',
                                        train=False,
                                        download=True,
                                        transform=trans_cifar)
        if args.iid:
            dict_users = cifar_iid(dataset_train, args.num_users)
        else:
            exit('Error: only consider IID setting in CIFAR10')
    else:
        exit('Error: unrecognized dataset')
    return dataset_train, dataset_test, dict_users
Exemplo n.º 4
0
    if args.dataset == 'mnist':
        trans_mnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
        dataset_train = datasets.MNIST('data/mnist/', train=True, download=True, transform=trans_mnist)
        dataset_test = datasets.MNIST('data/mnist/', train=False, download=True, transform=trans_mnist)
        # sample users
        if args.iid:
            dict_users = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users = mnist_noniid(dataset_train, args.num_users)
    elif args.dataset == 'cifar':
        trans_cifar_train = transforms.Compose([transforms.RandomCrop(32,padding=4),transforms.RandomHorizontalFlip(),transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])
        trans_cifar_test = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])
        dataset_train = datasets.CIFAR10('data/cifar', train=True, download=True, transform=trans_cifar_train)
        dataset_test = datasets.CIFAR10('data/cifar', train=False, download=True, transform=trans_cifar_test)
        if args.iid:
            dict_users = cifar_iid(dataset_train, args.num_users)
        else:
            exit('Error: only consider IID setting in CIFAR10')
    else:
        exit('Error: unrecognized dataset')
    img_size = dataset_train[0][0].shape

    # build model
    if args.model == 'cnn' and args.dataset == 'cifar':
        net_glob = customCNNCifar(args=args).to(args.device)
        net_glob1 = customCNNCifar(args=args).to(args.device)
        net_glob5 = customCNNCifar(args=args).to(args.device)
        net_glob10 = customCNNCifar(args=args).to(args.device)
        net_glob15 = customCNNCifar(args=args).to(args.device)
        net_glob20 = customCNNCifar(args=args).to(args.device)
        net_glob25 = customCNNCifar(args=args).to(args.device)
Exemplo n.º 5
0
                                      args.num_users)  #utils/sampling.py
    elif args.dataset == 'cifar':
        trans_cifar = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        dataset_train = datasets.CIFAR10('./data/cifar',
                                         train=True,
                                         download=True,
                                         transform=trans_cifar)
        dataset_test = datasets.CIFAR10('./data/cifar',
                                        train=False,
                                        download=True,
                                        transform=trans_cifar)
        if args.iid:
            dict_users = cifar_iid(dataset_train,
                                   args.num_users)  #utils/sampling.py
        else:
            exit('Error: only consider IID setting in CIFAR10')
    else:
        exit('Error: unrecognized dataset')
    img_size = dataset_train[0][0].shape

    # build model
    if args.model == 'cnn' and args.dataset == 'cifar':
        net_glob = CNNCifar(args=args).to(args.device)
    elif args.model == 'cnn' and args.dataset == 'mnist':
        net_glob = CNNMnist(args=args).to(args.device)
    elif args.model == 'mlp':
        len_in = 1
        for x in img_size:
            len_in *= x
Exemplo n.º 6
0
def main():
    # parse args
    args = args_parser()
    args.device = torch.device('cuda:{}'.format(
        args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')

    # load dataset and split users
    if args.dataset == 'mnist':
        trans_mnist = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307, ), (0.3081, ))
        ])
        dataset_train = datasets.MNIST('../data/mnist/',
                                       train=True,
                                       download=True,
                                       transform=trans_mnist)
        dataset_test = datasets.MNIST('../data/mnist/',
                                      train=False,
                                      download=True,
                                      transform=trans_mnist)
        print("type of test dataset", type(dataset_test))
        # sample users
        if args.iid:
            dict_users = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users, dict_labels_counter = mnist_noniid(
                dataset_train, args.num_users)
            dict_users_2, dict_labels_counter_2 = dict_users, dict_labels_counter
            #dict_users, dict_labels_counter = mnist_noniid_unequal(dataset_train, args.num_users)
    elif args.dataset == 'cifar':
        trans_cifar = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        dataset_train = datasets.CIFAR10('../data/cifar',
                                         train=True,
                                         download=True,
                                         transform=trans_cifar)
        dataset_test = datasets.CIFAR10('../data/cifar',
                                        train=False,
                                        download=True,
                                        transform=trans_cifar)
        if args.iid:
            dict_users = cifar_iid(dataset_train, args.num_users)
        else:
            exit('Error: only consider IID setting in CIFAR10')
    else:
        exit('Error: unrecognized dataset')
    img_size = dataset_train[0][0].shape

    # build model
    if args.model == 'cnn' and args.dataset == 'cifar':
        net_glob = CNNCifar(args=args).to(args.device)
        net_glob_2 = CNNCifar(args=args).to(args.device)
    elif args.model == 'cnn' and args.dataset == 'mnist':
        net_glob = CNNMnist(args=args).to(args.device)
        net_glob_2 = CNNMnist(args=args).to(args.device)
    elif args.model == 'mlp':
        len_in = 1
        for x in img_size:
            len_in *= x
        net_glob = MLP(dim_in=len_in, dim_hidden=200,
                       dim_out=args.num_classes).to(args.device)
    else:
        exit('Error: unrecognized model')

    #print(net_glob)

    #net_glob.train()

    acc_test, loss_test = test_img(net_glob, dataset_test, args)
    print("val test finished")
    print("{:.2f}".format(acc_test))
    temp = net_glob

    #net_glob_2 = net_glob
    temp_2 = net_glob_2

    # copy weights
    w_glob = net_glob.state_dict()

    # training
    loss_train = []
    cv_loss, cv_acc = [], []
    val_loss_pre, counter = 0, 0
    net_best = None
    best_loss = None
    val_acc_list, net_list = [], []

    Loss_local_each_global_total = []

    test_ds, valid_ds = torch.utils.data.random_split(dataset_test,
                                                      (9500, 500))
    loss_workers_total = np.zeros(shape=(args.num_users, args.epochs))
    label_workers = {
        i: np.array([], dtype='int64')
        for i in range(args.num_users)
    }

    workers_percent = []
    workers_count = 0
    acc_test_global, loss_test_global = test_img(x, valid_ds, args)
    selected_users_index = []

    for idx in range(args.num_users):
        # print("train started")
        local = LocalUpdate(args=args,
                            dataset=dataset_train,
                            idxs=dict_users[idx])
        w, loss = local.train(net=copy.deepcopy(net_glob).to(args.device))
        # print(w)
        # print("train completed")

        # temp = FedAvg(w)
        temp.load_state_dict(w)
        temp.eval()
        acc_test_local, loss_test_local = test_img(temp, valid_ds, args)
        loss_workers_total[idx, iter] = acc_test_local

        if workers_count >= (args.num_users / 2):
            break
        elif acc_test_local >= (0.7 * acc_test_global):
            selected_users_index.append(idx)

    for iter in range(args.epochs):
        print("round started")
        Loss_local_each_global = []
        loss_workers = np.zeros((args.num_users, args.epochs))
        w_locals, loss_locals = [], []
        m = max(int(args.frac * args.num_users), 1)
        #idxs_users = np.random.choice(range(args.num_users), m, replace=False)

        #if iter % 5 == 0:
        # Minoo
        x = net_glob
        x.eval()

        Loss_local_each_global_total.append(acc_test_global)

        for idx in selected_users_index:
            #print("train started")
            local = LocalUpdate(args=args,
                                dataset=dataset_train,
                                idxs=dict_users[idx])
            w, loss = local.train(net=copy.deepcopy(net_glob).to(args.device))
            #print(w)
            #print("train completed")

            #temp = FedAvg(w)
            temp.load_state_dict(w)
            temp.eval()
            acc_test_local, loss_test_local = test_img(temp, valid_ds, args)
            loss_workers_total[idx, iter] = acc_test_local

            if workers_count >= (args.num_users / 2):
                break
            elif acc_test_local >= (0.7 * acc_test_global):
                w_locals.append(copy.deepcopy(w))
                loss_locals.append(copy.deepcopy(loss))
                print("Update Received")
                workers_count += 1

        # update global weights
        w_glob = FedAvg(w_locals)

        # copy weight to net_glob
        net_glob.load_state_dict(w_glob)

        print("round completed")
        loss_avg = sum(loss_locals) / len(loss_locals)
        print('Round {:3d}, Average loss {:.3f}'.format(iter, loss_avg))
        loss_train.append(loss_avg)
        workers_percent.append(workers_count)

    # plot loss curve
    plt.figure()
    plt.plot(range(len(workers_percent)), workers_percent)
    plt.ylabel('train_loss')
    plt.savefig(
        './save/Newfed_WorkersPercent_0916_{}_{}_{}_C{}_iid{}.png'.format(
            args.dataset, args.model, args.epochs, args.frac, args.iid))
    # print(loss_workers_total)

    # plot loss curve
    # plt.figure()
    # plt.plot(range(len(loss_train)), loss_train)
    # plt.ylabel('train_loss')
    # plt.savefig('./save/Newfed_0916_{}_{}_{}_C{}_iid{}.png'.format(args.dataset, args.model, args.epochs, args.frac, args.iid))
    #

    plt.figure()
    for i in range(args.num_users):
        plot = plt.plot(range(len(loss_workers_total[i, :])),
                        loss_workers_total[i, :],
                        label="Worker {}".format(i))
    plot5 = plt.plot(range(len(Loss_local_each_global_total)),
                     Loss_local_each_global_total,
                     color='000000',
                     label="Global")
    plt.legend(loc='best')
    plt.ylabel('Small Test Set Accuracy of workers')
    plt.xlabel('Number of Rounds')
    plt.savefig(
        './save/NewFed_2workers_Acc_0916_{}_{}_{}_C{}_iid{}.png'.format(
            args.dataset, args.model, args.epochs, args.frac, args.iid))

    # plt.figure()
    # bins = np.linspace(0, 9, 3)
    # a = dict_labels_counter[:, 0].ravel()
    # print(type(a))
    # b = dict_labels_counter[:, 1].ravel()
    # x_labels = ['0', '1', '2', '3','4','5','6','7','8','9']
    # # Set plot parameters
    # fig, ax = plt.subplots()
    # width = 0.1  # width of bar
    # x = np.arange(10)
    # ax.bar(x, dict_labels_counter[:, 0], width, color='#000080', label='Worker 1')
    # ax.bar(x + width, dict_labels_counter[:, 1], width, color='#73C2FB', label='Worker 2')
    # ax.bar(x + 2*width, dict_labels_counter[:, 2], width, color='#ff0000', label='Worker 3')
    # ax.bar(x + 3*width, dict_labels_counter[:, 3], width, color='#32CD32', label='Worker 4')
    # ax.set_ylabel('Number of Labels')
    # ax.set_xticks(x + width + width / 2)
    # ax.set_xticklabels(x_labels)
    # ax.set_xlabel('Labels')
    # ax.legend()
    # plt.grid(True, 'major', 'y', ls='--', lw=.5, c='k', alpha=.3)
    # fig.tight_layout()
    # plt.savefig(
    #     './save/Newfed_2workersLabels_0916_{}_{}_{}_C{}_iid{}.png'.format(args.dataset, args.model, args.epochs, args.frac,
    #                                                                args.iid))

    # testing
    print("testing started")
    net_glob.eval()
    print("train test started")
    acc_train_final, loss_train_final = test_img(net_glob, dataset_train, args)
    print("train test finished")
    acc_test_final, loss_test_final = test_img(net_glob, dataset_test, args)
    print("val test finished")
    #print("Training accuracy: {:.2f}".format(acc_train))
    #print("Testing accuracy: {:.2f}".format(acc_test))
    print("{:.2f}".format(acc_test_final))
    #print("{:.2f".format(Loss_local_each_worker))

    # training
    w_glob_2 = net_glob_2.state_dict()

    loss_train_2 = []
    cv_loss_2, cv_acc_2 = [], []
    val_loss_pre_2, counter_2 = 0, 0
    net_best_2 = None
    best_loss_2 = None
    val_acc_list_2, net_list_2 = [], []

    Loss_local_each_global_total_2 = []

    loss_workers_total_2 = np.zeros(shape=(args.num_users, args.epochs))
    label_workers_2 = {
        i: np.array([], dtype='int64')
        for i in range(args.num_users)
    }

    for iter in range(args.epochs):
        print("round started")
        Loss_local_each_global_2 = []
        loss_workers_2 = np.zeros((args.num_users, args.epochs))
        w_locals_2, loss_locals_2 = [], []
        m_2 = max(int(args.frac * args.num_users), 1)
        idxs_users_2 = np.random.choice(range(args.num_users),
                                        m_2,
                                        replace=False)

        # Minoo
        x_2 = net_glob_2
        x_2.eval()
        acc_test_global_2, loss_test_global_2 = test_img(x_2, valid_ds, args)
        Loss_local_each_global_total_2.append(acc_test_global_2)

        for idx in idxs_users_2:
            #print("train started")
            local_2 = LocalUpdate(args=args,
                                  dataset=dataset_train,
                                  idxs=dict_users_2[idx])
            w_2, loss_2 = local_2.train(
                net=copy.deepcopy(net_glob_2).to(args.device))
            #print(w)
            #print("train completed")
            w_locals_2.append(copy.deepcopy(w_2))
            loss_locals_2.append(copy.deepcopy(loss_2))
            #temp = FedAvg(w)
            temp_2.load_state_dict(w_2)
            temp_2.eval()
            acc_test_local_2, loss_test_local_2 = test_img(
                temp_2, valid_ds, args)
            loss_workers_total_2[idx, iter] = acc_test_local_2

        # update global weights
        w_glob_2 = FedAvg(w_locals_2)

        # copy weight to net_glob
        net_glob_2.load_state_dict(w_glob_2)

        print("round completed")
        loss_avg_2 = sum(loss_locals_2) / len(loss_locals_2)
        print('Round {:3d}, Average loss {:.3f}'.format(iter, loss_avg_2))
        loss_train_2.append(loss_avg_2)
        print("round completed")

        # plot loss curve
    plt.figure()
    plt.plot(range(len(loss_train_2)),
             loss_train_2,
             color='#000000',
             label="Main FL")
    plt.plot(range(len(loss_train)),
             loss_train,
             color='#ff0000',
             label="Centralized Algorithm")
    plt.ylabel('train_loss')
    plt.savefig('./save/main_fed_0916_{}_{}_{}_C{}_iid{}.png'.format(
        args.dataset, args.model, args.epochs, args.frac, args.iid))
    # print(loss_workers_total)

    plt.figure()
    for i in range(args.num_users):
        plot = plt.plot(range(len(loss_workers_total_2[i, :])),
                        loss_workers_total_2[i, :],
                        label="Worker {}".format(i))
    plot5 = plt.plot(range(len(Loss_local_each_global_total_2)),
                     Loss_local_each_global_total_2,
                     color='000000',
                     label="Global")
    plt.legend(loc='best')
    plt.ylabel('Small Test Set Accuracy of workers')
    plt.xlabel('Number of Rounds')
    plt.savefig('./save/mainfed_Acc_0916_{}_{}_{}_C{}_iid{}.png'.format(
        args.dataset, args.model, args.epochs, args.frac, args.iid))

    # plt.figure()
    # bins = np.linspace(0, 9, 3)
    # a = dict_labels_counter_2[:, 0].ravel()
    # print(type(a))
    # b = dict_labels_counter_2[:, 1].ravel()
    # x_labels = ['0', '1', '2', '3','4','5','6','7','8','9']
    # # Set plot parameters
    # fig, ax = plt.subplots()
    # width = 0.1  # width of bar
    # x = np.arange(10)
    # ax.bar(x, dict_labels_counter_2[:, 0], width, color='#000080', label='Worker 1')
    # ax.bar(x + width, dict_labels_counter_2[:, 1], width, color='#73C2FB', label='Worker 2')
    # ax.bar(x + 2*width, dict_labels_counter_2[:, 2], width, color='#ff0000', label='Worker 3')
    # ax.bar(x + 3*width, dict_labels_counter_2[:, 3], width, color='#32CD32', label='Worker 4')
    # ax.set_ylabel('Number of Labels')
    # ax.set_xticks(x + width + width / 2)
    # ax.set_xticklabels(x_labels)
    # ax.set_xlabel('Labels')
    # ax.legend()
    # plt.grid(True, 'major', 'y', ls='--', lw=.5, c='k', alpha=.3)
    # fig.tight_layout()
    # plt.savefig(
    #     './save/main_fed_2workersLabels_0916_{}_{}_{}_C{}_iid{}.png'.format(args.dataset, args.model, args.epochs, args.frac,
    #                                                                args.iid))

    # testing
    print("testing started")
    net_glob.eval()
    print("train test started")
    acc_train_final, loss_train_final = test_img(net_glob, dataset_train, args)
    print("train test finished")
    acc_test_final, loss_test_final = test_img(net_glob, dataset_test, args)
    print("val test finished")
    #print("Training accuracy: {:.2f}".format(acc_train))
    #print("Testing accuracy: {:.2f}".format(acc_test))
    print("{:.2f}".format(acc_test_final))
    #print("{:.2f".format(Loss_local_each_worker))

    return loss_test_final, loss_train_final
Exemplo n.º 7
0
def main():

    manualSeed = 1

    np.random.seed(manualSeed)
    random.seed(manualSeed)
    torch.manual_seed(manualSeed)
    # if you are suing GPU
    torch.cuda.manual_seed(manualSeed)
    torch.cuda.manual_seed_all(manualSeed)

    torch.backends.cudnn.enabled = False
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

    # parse args
    args = args_parser()
    args.device = torch.device('cuda:{}'.format(
        args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')

    # load dataset and split users
    if args.dataset == 'mnist':
        trans_mnist = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307, ), (0.3081, ))
        ])
        dataset_train = datasets.MNIST('../data/mnist/',
                                       train=True,
                                       download=True,
                                       transform=trans_mnist)
        dataset_test = datasets.MNIST('../data/mnist/',
                                      train=False,
                                      download=True,
                                      transform=trans_mnist)
        # sample users
        if args.iid:
            dict_users_DCFL = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users_DCFL, dict_labels_counter = mnist_noniid(
                dataset_train, args.num_users)
            dict_users_mainFL, dict_labels_counter_mainFL = dict_users_DCFL, dict_labels_counter
    elif args.dataset == 'cifar':
        trans_cifar = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        dataset_train = datasets.CIFAR10('../data/cifar',
                                         train=True,
                                         download=True,
                                         transform=trans_cifar)
        dataset_test = datasets.CIFAR10('../data/cifar',
                                        train=False,
                                        download=True,
                                        transform=trans_cifar)
        if args.iid:
            dict_users_DCFL = cifar_iid(dataset_train, args.num_users)
            dict_users_mainFL = dict_users_DCFL
            dict_labels_counter_mainFL = dict()
            dict_labels_counter = dict()
        else:
            dict_users_DCFL, dict_labels_counter = cifar_noniid(
                dataset_train, args.num_users)
            dict_users_mainFL, dict_labels_counter_mainFL = dict_users_DCFL, dict_labels_counter
    elif args.dataset == 'fmnist':
        trans_fmnist = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307, ), (0.3081, ))
        ])
        dataset_train = datasets.FashionMNIST('../data/fmnist',
                                              train=True,
                                              download=True,
                                              transform=trans_fmnist)
        dataset_test = datasets.FashionMNIST('../data/fmnist',
                                             train=False,
                                             download=True,
                                             transform=trans_fmnist)
        if args.iid:
            print("iid")
            dict_users_DCFL = mnist_iid(dataset_train, args.num_users)
        else:
            print("non iid")
            dict_users_DCFL, dict_labels_counter = mnist_noniid(
                dataset_train, args.num_users)
            dict_users_mainFL, dict_labels_counter_mainFL = dict_users_DCFL, dict_labels_counter
    else:
        exit('Error: unrecognized dataset')

    img_size = dataset_train[0][0].shape

    # Small shared dataset
    test_ds, valid_ds_before = torch.utils.data.random_split(
        dataset_test, (9500, 500))
    small_shared_dataset = create_shared_dataset(valid_ds_before, 200)

    optimal_delay = 1.0

    # Start process for each fraction of c
    for c_counter in range(3, 3 + 1, 2):
        if args.model == 'cnn' and args.dataset == 'cifar':
            net_glob = CNNCifar(args=args).to(args.device)
            # net_glob_mainFL = copy.deepcopy(net_glob)
        elif args.model == 'cnn' and args.dataset == 'mnist':
            net_glob = CNNMnist(args=args).to(args.device)
            # net_glob_mainFL = copy.deepcopy(net_glob)
        elif args.model == 'cnn' and args.dataset == 'fmnist':
            net_glob = CNNFashion_Mnist(args=args).to(args.device)
            # net_glob_mainFL = copy.deepcopy(net_glob)
        elif args.model == 'mlp':
            len_in = 1
            for x in img_size:
                len_in *= x
            net_glob = MLP(dim_in=len_in,
                           dim_hidden=200,
                           dim_out=args.num_classes).to(args.device)
        else:
            exit('Error: unrecognized model')

        # Saving data
        data_Global_main = {
            "C": [],
            "Round": [],
            "Average Loss Train": [],
            "SDS Loss": [],
            "SDS Accuracy": [],
            "Workers Number": [],
            "Large Test Loss": [],
            "Large Test Accuracy": [],
            "Communication Cost": []
        }
        Final_LargeDataSetTest_MainFL = {
            "C": [],
            "Test Accuracy": [],
            "Test Loss": [],
            "Train Loss": [],
            "Train Accuracy": [],
            "Total Rounds": [],
            "Communication Cost": []
        }

        data_Global_DCFL = {
            "C": [],
            "Round": [],
            "Average Loss Train": [],
            "SDS Loss": [],
            "SDS Accuracy": [],
            "Workers Number": [],
            "Large Test Loss": [],
            "Large Test Accuracy": [],
            "Communication Cost": []
        }
        Final_LargeDataSetTest_DCFL = {
            "C": [],
            "Test Accuracy": [],
            "Test Loss": [],
            "Train Loss": [],
            "Train Accuracy": [],
            "Total Rounds": [],
            "Communication Cost": []
        }

        data_Global_G1 = {
            "C": [],
            "Round": [],
            "Average Loss Train": [],
            "SDS Loss": [],
            "SDS Accuracy": [],
            "Workers Number": [],
            "Large Test Loss": [],
            "Large Test Accuracy": [],
            "Communication Cost": []
        }
        Final_LargeDataSetTest_G1 = {
            "C": [],
            "Test Accuracy": [],
            "Test Loss": [],
            "Train Loss": [],
            "Train Accuracy": [],
            "Total Rounds": [],
            "Communication Cost": []
        }

        data_Global_G2 = {
            "C": [],
            "Round": [],
            "Average Loss Train": [],
            "SDS Loss": [],
            "SDS Accuracy": [],
            "Workers Number": [],
            "Large Test Loss": [],
            "Large Test Accuracy": [],
            "Communication Cost": []
        }
        Final_LargeDataSetTest_G2 = {
            "C": [],
            "Test Accuracy": [],
            "Test Loss": [],
            "Train Loss": [],
            "Train Accuracy": [],
            "Total Rounds": [],
            "Communication Cost": []
        }

        data_Global_Muhammed = {
            "C": [],
            "Round": [],
            "Average Loss Train": [],
            "SDS Loss": [],
            "SDS Accuracy": [],
            "Workers Number": [],
            "Large Test Loss": [],
            "Large Test Accuracy": [],
            "Communication Cost": []
        }
        Final_LargeDataSetTest_Muhammed = {
            "C": [],
            "Test Accuracy": [],
            "Test Loss": [],
            "Train Loss": [],
            "Train Accuracy": [],
            "Total Rounds": [],
            "Communication Cost": []
        }

        data_Global_Cho = {
            "C": [],
            "Round": [],
            "Average Loss Train": [],
            "SDS Loss": [],
            "SDS Accuracy": [],
            "Workers Number": [],
            "Large Test Loss": [],
            "Large Test Accuracy": [],
            "Communication Cost": []
        }
        Final_LargeDataSetTest_Cho = {
            "C": [],
            "Test Accuracy": [],
            "Test Loss": [],
            "Train Loss": [],
            "Train Accuracy": [],
            "Total Rounds": [],
            "Communication Cost": []
        }

        net_glob.train()
        net_glob_mainFL = copy.deepcopy(net_glob)
        net_glob_G1 = copy.deepcopy(net_glob)
        net_glob_G2 = copy.deepcopy(net_glob)
        cost = np.random.rand(args.num_users)

        R_G1 = 5
        args.frac = (c_counter / 10)

        # Main FL
        loss_main, dict_workers_index, Final_LargeDataSetTest_MainFL_temp, data_Global_main_temp = mainFl(
            net_glob_mainFL, dict_users_mainFL, dict_labels_counter_mainFL,
            args, cost, dataset_train, dataset_test, small_shared_dataset)

        Final_LargeDataSetTest_MainFL = merge(
            Final_LargeDataSetTest_MainFL, Final_LargeDataSetTest_MainFL_temp)
        data_Global_main = merge(data_Global_main, data_Global_main_temp)

        # with open(os.path.join(OUT_DIR, f"dict_users_mainFL-C-{args.frac}-{args.dataset}.pkl"), 'wb') as file:
        #     pickle.dump(dict_users_mainFL, file)

        # with open(os.path.join(OUT_DIR, f"dict_users_mainFL-C-{args.frac}-{args.dataset}.pkl"), 'rb') as file:
        #     dict_users_mainFL = pickle.load(file)

        # with open(os.path.join(OUT_DIR, f"workers_index-C-{args.frac}-{args.dataset}.pkl"), 'wb') as file:
        #     pickle.dump(dict_workers_index, file)

        # with open(os.path.join(OUT_DIR, f"cost-C-{args.frac}-{args.dataset}.pkl"), 'wb') as file:
        #     pickle.dump(cost, file)

        # with open(os.path.join(OUT_DIR, f"cost-C-{args.frac}-{args.dataset}.pkl"), 'rb') as file:
        #     cost = pickle.load(file)

        # print(cost)

        # with open(os.path.join(OUT_DIR, f"GoalLoss-C-{args.frac}-{args.dataset}.pkl"), 'wb') as file:
        #     pickle.dump(loss_main, file)

        date = datetime.now()
        _dir = os.path.join(OUT_DIR, str(date.date()))
        if not os.path.exists(_dir):
            os.makedirs(_dir)
        save_time = time.strftime("%Y%m%d-%H%M%S")

        Final_LargeDataSetTest_MainFL = pd.DataFrame.from_dict(
            Final_LargeDataSetTest_MainFL)
        data_Global_main = pd.DataFrame.from_dict(data_Global_main)
        Final_LargeDataSetTest_MainFL.to_csv(
            os.path.join(
                _dir,
                f"{save_time}-{args.dataset}-Final_LargeDataSetTest_MainFL.csv"
            ))
        data_Global_main.to_csv(
            os.path.join(_dir,
                         f"{save_time}-{args.dataset}-data_Global_main.csv"))

        # Proposed G1
        Final_LargeDataSetTest_G1_temp, data_Global_G1_temp = Proposed_G1(
            net_glob_G1, dict_workers_index, dict_users_DCFL,
            dict_labels_counter_mainFL, args, cost, dataset_train,
            dataset_test, small_shared_dataset, loss_main, R_G1, optimal_delay)
        Final_LargeDataSetTest_G1 = merge(Final_LargeDataSetTest_G1,
                                          Final_LargeDataSetTest_G1_temp)
        data_Global_G1 = merge(data_Global_G1, data_Global_G1_temp)

        Final_LargeDataSetTest_G1 = pd.DataFrame.from_dict(
            Final_LargeDataSetTest_G1)
        data_Global_G1 = pd.DataFrame.from_dict(data_Global_G1)
        Final_LargeDataSetTest_G1.to_csv(
            os.path.join(
                _dir,
                f"{save_time}-{args.dataset}-Final_LargeDataSetTest_G1.csv"))
        data_Global_G1.to_csv(
            os.path.join(_dir,
                         f"{save_time}-{args.dataset}-data_Global_G1.csv"))

        print("G1 alg is done")
Exemplo n.º 8
0
def main():
    # parse args
    args = args_parser()
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    dataPath = args.datasetPath

    # random seed
    np.random.seed(args.seed)
    cudnn.benchmark = False
    cudnn.deterministic = True
    torch.manual_seed(args.seed)
    cudnn.enabled = True
    torch.cuda.manual_seed(args.seed)

    # load dataset and split users
    if args.dataset == 'cifar10':
        _CIFAR_TRAIN_TRANSFORMS = [
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465),
                                 (0.2023, 0.1994, 0.2010)),
        ]
        dataset_train = datasets.CIFAR10(
            dataPath,
            train=True,
            download=True,
            transform=transforms.Compose(_CIFAR_TRAIN_TRANSFORMS))

        _CIFAR_TEST_TRANSFORMS = [
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465),
                                 (0.2023, 0.1994, 0.2010)),
        ]
        dataset_test = datasets.CIFAR10(
            dataPath,
            train=False,
            transform=transforms.Compose(_CIFAR_TEST_TRANSFORMS))

        if args.iid == 0:  # IID
            dict_users = cifar_iid(dataset_train, args.num_users)
        elif args.iid == 2:  # non-IID
            dict_users = cifar_noniid_2(dataset_train, args.num_users)
        else:
            exit('Error: unrecognized class')

    elif args.dataset == 'emnist':
        _MNIST_TRAIN_TRANSFORMS = _MNIST_TEST_TRANSFORMS = [
            transforms.ToTensor(),
            transforms.ToPILImage(),
            transforms.Pad(2),
            transforms.ToTensor(),
            transforms.Normalize((0.1307, ), (0.3081, ))
        ]
        dataset_train = datasets.EMNIST(
            dataPath,
            train=True,
            download=True,
            transform=transforms.Compose(_MNIST_TRAIN_TRANSFORMS),
            split='letters')
        dataset_test = datasets.EMNIST(
            dataPath,
            train=False,
            download=True,
            transform=transforms.Compose(_MNIST_TEST_TRANSFORMS),
            split='letters')

        dict_users = femnist_star(dataset_train, args.num_users)

    elif args.dataset == 'cifar100':
        _CIFAR_TRAIN_TRANSFORMS = [
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465),
                                 (0.2023, 0.1994, 0.2010)),
        ]
        dataset_train = datasets.CIFAR100(
            dataPath,
            train=True,
            download=True,
            transform=transforms.Compose(_CIFAR_TRAIN_TRANSFORMS))

        _CIFAR_TEST_TRANSFORMS = [
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465),
                                 (0.2023, 0.1994, 0.2010)),
        ]
        dataset_test = datasets.CIFAR100(
            dataPath,
            train=False,
            transform=transforms.Compose(_CIFAR_TEST_TRANSFORMS))
        if args.iid == 0:  # IID
            dict_users = cifar_100_iid(dataset_train, args.num_users)
        elif args.iid == 2:  # non-IID
            dict_users = cifar_100_noniid(dataset_train, args.num_users)
    else:
        exit('Error: unrecognized dataset')

    # build model
    if args.dataset == 'cifar10':
        if args.model == "CNNStd5":
            net_glob = CNNCifarStd5().cuda()
        else:
            exit('Error: unrecognized model')
    elif args.dataset == 'emnist':
        if args.model == "CNNStd5":
            net_glob = CNNEmnistStd5().cuda()
        else:
            exit('Error: unrecognized model')
    elif args.dataset == 'cifar100':
        if args.model == "CNNStd5":
            net_glob = CNNCifar100Std5().cuda()
        else:
            exit('Error: unrecognized model')
    else:
        exit('Error: unrecognized model')

    print('Number of model parameters: {}'.format(
        sum([p.data.nelement() for p in net_glob.parameters()])))

    net_glob.train()

    learning_rate = args.lr
    test_acc = []
    avg_loss = []

    # Train
    for iter in range(args.epochs):

        m = max(int(args.frac * args.num_users), 1)
        idxs_users = np.random.choice(range(args.num_users), m, replace=False)
        w_locals, loss_locals = [], []
        for i, idx in enumerate(idxs_users):
            print('user: {:d}'.format(idx))
            local = LocalUpdate(args=args,
                                dataset=dataset_train,
                                idxs=dict_users[idx])
            w, loss = local.train(model=copy.deepcopy(net_glob).cuda(),
                                  lr=learning_rate)

            w_locals.append(copy.deepcopy(w))
            loss_locals.append(copy.deepcopy(loss))

        # update global weights
        w_glob = FedAvg(w_locals)

        # copy weight to net_glob
        net_glob.load_state_dict(w_glob)

        # print loss
        loss_avg = sum(loss_locals) / len(loss_locals)
        print('Round {:3d}, Average loss {:.6f}'.format(iter, loss_avg))

        acc_test, _ = test_img(net_glob.cuda(), dataset_test, args)
        print("test accuracy: {:.4f}".format(acc_test))
        test_acc.append(acc_test)

        avg_loss.append(loss_avg)

        learning_rate = adjust_learning_rate(learning_rate, args.lr_drop)

    filename = './accuracy-' + str(args.dataset) + '-iid' + str(args.iid) + '-' + str(args.epochs) + '-seed' \
               + str(args.seed) + '-' + str(args.loss_type) + '-beta' + str(args.beta) + '-mu' + str(args.mu)
    save_result(test_acc, avg_loss, filename)
Exemplo n.º 9
0
def modelBuild():
    """
    Build the basic training network and return the related args.
    """
    # build model
    args = args_parser()
    args.device = torch.device('cuda:{}'.format(
        args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')

    # load dataset and split users
    if args.dataset == 'mnist':
        trans_mnist = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307, ), (0.3081, ))
        ])
        dataset_train = datasets.MNIST('../data/mnist/',
                                       train=True,
                                       download=True,
                                       transform=trans_mnist)
        dataset_test = datasets.MNIST('../data/mnist/',
                                      train=False,
                                      download=True,
                                      transform=trans_mnist)
        # sample users
        if args.iid:
            # allocate the dataset index to users
            dict_users = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users = mnist_noniid(dataset_train, args.num_users)
    elif args.dataset == 'cifar':
        trans_cifar = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        dataset_train = datasets.CIFAR10('../data/cifar',
                                         train=True,
                                         download=True,
                                         transform=trans_cifar)
        dataset_test = datasets.CIFAR10('../data/cifar',
                                        train=False,
                                        download=True,
                                        transform=trans_cifar)
        if args.iid:
            dict_users = cifar_iid(dataset_train, args.num_users)
        else:
            exit('Error: only consider IID setting in CIFAR10')
    else:
        exit('Error: unrecognized dataset')

    print("The para of iid is " + str(args.iid))

    img_size = dataset_train[0][0].shape
    if args.model == 'cnn' and args.dataset == 'cifar':
        net_glob = CNNCifar(args=args).to(args.device)
    elif args.model == 'cnn' and args.dataset == 'mnist':
        net_glob = CNNMnist(args=args).to(args.device)
    elif args.model == 'mlp':
        len_in = 1
        for x in img_size:
            len_in *= x
        net_glob = MLP(dim_in=len_in, dim_hidden=200,
                       dim_out=args.num_classes).to(args.device)
    else:
        exit('Error: unrecognized model')

    print("********************************")
    print(net_glob)
    print("********************************")

    return net_glob, args, dataset_train, dataset_test, dict_users
Exemplo n.º 10
0
def main():
    # parse args
    args = args_parser()
    args.device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')

    # load dataset and split users
    if args.dataset == 'mnist':
        trans_mnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
        dataset_train = datasets.MNIST('../data/mnist/', train=True, download=True, transform=trans_mnist)
        dataset_test = datasets.MNIST('../data/mnist/', train=False, download=True, transform=trans_mnist)
        # sample users
        if args.iid:
            dict_users = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users, dict_labels_counter = mnist_noniid(dataset_train, args.num_users)
            dict_users_mainFL, dict_labels_counter_mainFL = dict_users, dict_labels_counter
    elif args.dataset == 'cifar':
        trans_cifar = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        dataset_train = datasets.CIFAR10('../data/cifar', train=True, download=True, transform=trans_cifar)
        dataset_test = datasets.CIFAR10('../data/cifar', train=False, download=True, transform=trans_cifar)
        if args.iid:
            dict_users = cifar_iid(dataset_train, args.num_users)
        else:
            dict_users, dict_labels_counter = cifar_noniid(dataset_train, args.num_users)
            dict_users_mainFL, dict_labels_counter_mainFL = dict_users, dict_labels_counter
    elif args.dataset == 'fmnist':
        trans_fmnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
        dataset_train = datasets.FashionMNIST('../data/fmnist', train=True, download=True, transform=trans_fmnist)
        dataset_test = datasets.FashionMNIST('../data/fmnist', train=False, download=True, transform=trans_fmnist)
        if args.iid:
            dict_users = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users, dict_labels_counter = mnist_noniid(dataset_train, args.num_users)
            dict_users_mainFL, dict_labels_counter_mainFL = dict_users, dict_labels_counter
    else:
        exit('Error: unrecognized dataset')


    img_size = dataset_train[0][0].shape

    acc_full_distributed = []
    acc_full_main = []
    loss_full_ditributed = []
    loss_full_main = []

    SD_acc_full_distributed = []
    SD_acc_full_main = []
    SD_loss_full_ditributed = []
    SD_loss_full_main = []

    workers_percent_full_distributed = []
    workers_percent_full_main = []
    variable_start = 0.1
    variable_end = 1.0
    while_counter = 0.1
    counter_array = []
    Accuracy_Fraction = []
    Workers_Fraction = []

    accuracy_fraction_each_round_newFL = 0
    workers_fraction_each_round_newFL = 0
    accuracy_fraction_each_round_mainFL = 0
    workers_fraction_each_round_mainFL = 0

    data_main = {}
    data_DCFL = {}
    data_Global_main = {"C": [], "Round":[], "Average Loss Train": [], "Average Loss Test": [], "Accuracy Test": [],
                        "Workers Number": [], "Large Test Loss":[], "Large Test Accuracy":[]}
    data_Global_DCFL = {"C": [], "Round":[], "Average Loss Train": [], "Average Loss Test": [], "Accuracy Test": [],
                        "Workers Number": [], "Large Test Loss":[], "Large Test Accuracy":[]}
    Final_LargeDataSetTest_DCFL = {"C":[], "Test Accuracy":[], "Test Loss":[], "Train Loss":[], "Train Accuracy":[],
                                   "Total Rounds":[]}
    Final_LargeDataSetTest_MainFL = {"C":[], "Test Accuracy": [], "Test Loss": [], "Train Loss": [], "Train Accuracy":[]}



    # build model
    args.frac = variable_start

    test_ds, valid_ds_before = torch.utils.data.random_split(dataset_test, (9500, 500))
    valid_ds = create_shared_dataset(valid_ds_before, 200)

    #while variable_start <= variable_end:
    for c_counter in range(1, 11, 3):
        if args.model == 'cnn' and args.dataset == 'cifar':
            net_glob = CNNCifar(args=args).to(args.device)
            net_glob_mainFL = copy.deepcopy(net_glob)
        elif args.model == 'cnn' and args.dataset == 'mnist':
            net_glob = CNNMnist(args=args).to(args.device)
            net_glob_mainFL = copy.deepcopy(net_glob)
        elif args.model == 'cnn' and args.dataset == 'fmnist':
            net_glob = CNNFashion_Mnist(args=args).to(args.device)
            net_glob_mainFL = copy.deepcopy(net_glob)
        elif args.model == 'mlp':
            len_in = 1
            for x in img_size:
                len_in *= x
            net_glob = MLP(dim_in=len_in, dim_hidden=200, dim_out=args.num_classes).to(args.device)
        else:
            exit('Error: unrecognized model')

        counter_array.append((c_counter/10))
        args.frac = (c_counter/10)

        ######saving index of workers
        dict_workers_index = defaultdict(list)


        #############Main FL

        w_glob_mainFL = net_glob_mainFL.state_dict()

        loss_train_mainFL = []
        # cv_loss_2, cv_acc_2 = [], []
        # val_loss_pre_2, counter_2 = 0, 0
        # net_best_2 = None
        # best_loss_2 = None
        # val_acc_list_2, net_list_2 = [], []

        Loss_local_each_global_total_mainFL = []
        Accuracy_local_each_global_total_mainFL = []

        loss_workers_total_mainFL = np.zeros(shape=(args.num_users, args.epochs))
        label_workers_mainFL = {i: np.array([], dtype='int64') for i in range(args.num_users)}

        validation_test_mainFed = []
        acc_test, loss_test = test_img(net_glob_mainFL, dataset_test, args)
        workers_participation_main_fd = np.zeros((args.num_users, args.epochs))
        workers_percent_main = []

        # for iter in range(args.epochs):
        net_glob_mainFL.eval()
        acc_test_final_mainFL, loss_test_final_mainFL = test_img(net_glob_mainFL, dataset_test, args)
        while_counter_mainFL = loss_test_final_mainFL
        iter_mainFL = 0

        workers_mainFL = []
        for i in range(args.num_users):
            workers_mainFL.append(i)

        temp_netglob_mainFL = net_glob_mainFL

        while iter_mainFL < (args.epochs/2):

            data_main['round_{}'.format(iter_mainFL)] = []
            # data_Global_main['round_{}'.format(iter)] = []
            # print("round started")
            Loss_local_each_global_mainFL = []
            loss_workers_mainFL = np.zeros((args.num_users, args.epochs))
            w_locals_mainFL, loss_locals_mainFL = [], []
            m_mainFL = max(int(args.frac * args.num_users), 1)
            idxs_users_mainFL = np.random.choice(range(args.num_users), m_mainFL, replace=False)
            list_of_random_workers = random.sample(workers_mainFL, m_mainFL)
            for i in range(len(list_of_random_workers)):
                dict_workers_index[iter_mainFL].append(list_of_random_workers[i])

            x_mainFL = net_glob_mainFL
            x_mainFL.eval()
            acc_test_global_mainFL, loss_test_global_mainFL = test_img(x_mainFL, valid_ds, args)
            Loss_local_each_global_total_mainFL.append(loss_test_global_mainFL)
            Accuracy_local_each_global_total_mainFL.append(acc_test_global_mainFL)
            SD_acc_full_main.append(acc_test_global_mainFL)
            SD_loss_full_main.append(loss_test_global_mainFL)

            workers_count_mainFL = 0
            temp_accuracy = np.zeros(1)
            temp_loss_test = np.zeros(1)
            temp_loss_train = np.zeros(1)
            for idx in list_of_random_workers:
                # print("train started")
                local_mainFL = LocalUpdate(args=args, dataset=dataset_train, idxs=dict_users_mainFL[idx])
                w_mainFL, loss_mainFL = local_mainFL.train(net=copy.deepcopy(net_glob_mainFL).to(args.device))
                # print(w)
                # print("train completed")
                w_locals_mainFL.append(copy.deepcopy(w_mainFL))
                loss_locals_mainFL.append(copy.deepcopy(loss_mainFL))
                # temp = FedAvg(w)
                temp_netglob_mainFL.load_state_dict(w_mainFL)
                temp_netglob_mainFL.eval()
                print(pnorm_2(temp_netglob_mainFL, 2))
                acc_test_local_mainFL, loss_test_local_mainFL = test_img(temp_netglob_mainFL, valid_ds, args)
                temp_accuracy[0] = acc_test_local_mainFL
                temp_loss_test[0] = loss_test_local_mainFL
                temp_loss_train[0] = loss_mainFL
                loss_workers_total_mainFL[idx, iter_mainFL] = acc_test_local_mainFL
                workers_participation_main_fd[idx][iter_mainFL] = 1
                workers_count_mainFL += 1
                data_main['round_{}'.format(iter_mainFL)].append({
                    'C': args.frac,
                    'User ID': idx,
                    # 'Local Update': copy.deepcopy(w_mainFL),
                    'Loss Train': temp_loss_train[0],
                    'Loss Test': temp_loss_test[0],
                    'Accuracy': temp_accuracy[0]
                })

            # update global weights
            w_glob_mainFL = FedAvg(w_locals_mainFL)

            # copy weight to net_glob
            net_glob_mainFL.load_state_dict(w_glob_mainFL)

            # print("round completed")
            loss_avg_mainFL = sum(loss_locals_mainFL) / len(loss_locals_mainFL)
            # print('Round {:3d}, Average loss {:.3f}'.format(iter, loss_avg_mainFL))
            loss_train_mainFL.append(loss_avg_mainFL)
            # print("round completed")

            acc_test_round_mainfed, loss_test_round_mainfed = test_img(net_glob_mainFL, dataset_test, args)
            validation_test_mainFed.append(acc_test_round_mainfed)
            workers_percent_main.append(workers_count_mainFL / args.num_users)

            # plot workers percent of participating
            print(iter_mainFL, " round main fl finished")

            acc_test_final_mainFL, loss_test_final_mainFL = test_img(net_glob_mainFL, dataset_test, args)
            while_counter_mainFL = loss_test_final_mainFL

            data_Global_main["Round"].append(iter_mainFL)
            data_Global_main["C"].append(args.frac)
            data_Global_main["Average Loss Train"].append(float(loss_avg_mainFL))
            data_Global_main["Average Loss Test"].append(float(loss_test_global_mainFL))
            data_Global_main["Accuracy Test"].append(float(acc_test_global_mainFL))
            data_Global_main["Workers Number"].append(float(workers_count_mainFL))
            data_Global_main["Large Test Loss"].append(float(loss_test_final_mainFL))
            data_Global_main["Large Test Accuracy"].append(float(acc_test_final_mainFL))

            iter_mainFL = iter_mainFL + 1

        workers_percent_final_mainFL = np.zeros(args.num_users)
        workers_name_mainFL = np.empty(args.num_users)
        for i in range(len(workers_participation_main_fd[:, 1])):
            workers_percent_final_mainFL[i] = sum(workers_participation_main_fd[i, :]) / args.epochs
            workers_name_mainFL[i] = i

        net_glob_mainFL.eval()
        # print("train test started")
        acc_train_final_main, loss_train_final_main = test_img(net_glob_mainFL, dataset_train, args)
        # print("train test finished")
        acc_test_final_main, loss_test_final_main = test_img(net_glob_mainFL, dataset_test, args)

        Final_LargeDataSetTest_MainFL["C"].append(args.frac)
        Final_LargeDataSetTest_MainFL["Test Loss"].append(float(loss_test_final_main))
        Final_LargeDataSetTest_MainFL["Test Accuracy"].append(float(acc_test_final_main))
        Final_LargeDataSetTest_MainFL["Train Loss"].append(float(loss_train_final_main))
        Final_LargeDataSetTest_MainFL["Train Accuracy"].append(float(acc_train_final_main))






        # copy weights
        w_glob = net_glob.state_dict()

        temp_after = copy.deepcopy(net_glob)
        temp_before = copy.deepcopy(net_glob)

        # training
        loss_train = []
        # cv_loss, cv_acc = [], []
        # val_loss_pre, counter = 0, 0
        # net_best = None
        # best_loss = None
        # val_acc_list, net_list = [], []

        Loss_local_each_global_total = []


        # valid_ds = create_shared_dataset(dataset_test, 500)
        loss_workers_total = np.zeros(shape=(args.num_users, args.epochs))
        label_workers = {i: np.array([], dtype='int64') for i in range(args.num_users)}

        workers_percent_dist = []
        validation_test_newFed = []
        workers_participation = np.zeros((args.num_users, args.epochs))
        workers = []
        for i in range(args.num_users):
            workers.append(i)

        counter_threshold_decrease = np.zeros(args.epochs)
        Global_Accuracy_Tracker = np.zeros(args.epochs)
        Global_Loss_Tracker = np.zeros(args.epochs)
        threshold = 0.5
        alpha = 0.5     ##decrease parameter
        beta = 0.1 ##delta accuracy controller
        gamma = 0.5  ##threshold decrease parameter


        Goal_Loss = float(loss_test_final_main)

        #for iter in range(args.epochs):

        net_glob.eval()
        acc_test_final, loss_test_final = test_img(net_glob, dataset_test, args)
        while_counter = float(loss_test_final)
        iter = 0

        total_rounds_dcfl = 0

        while (while_counter + 0.01) > Goal_Loss and iter <= args.epochs:

            data_DCFL['round_{}'.format(iter)] = []
            Loss_local_each_global = []
            loss_workers = np.zeros((args.num_users, args.epochs))
            w_locals, loss_locals = [], []
            m = max(int(args.frac * args.num_users), 1)
            idxs_users = np.random.choice(range(args.num_users), m, replace=False)
            counter_threshold = 0
            print(iter, " in dist FL started")
            #if iter % 5 == 0:

            x = copy.deepcopy(net_glob)
            x.eval()
            acc_test_global, loss_test_global = test_img(x, valid_ds, args)
            Loss_local_each_global_total.append(acc_test_global)
            Global_Accuracy_Tracker[iter] = acc_test_global
            Global_Loss_Tracker[iter] = loss_test_global
            if iter > 0 & (Global_Loss_Tracker[iter-1] - Global_Loss_Tracker[iter] <= beta):
                threshold = threshold - gamma
                if threshold == 0.0:
                    threshold = 1.0
                print("threshold decreased to", threshold)
            workers_count = 0

            SD_acc_full_distributed.append(acc_test_global)
            SD_loss_full_ditributed.append(loss_test_global)


            temp_w_locals = []
            temp_workers_loss = np.empty(args.num_users)
            temp_workers_accuracy = np.empty(args.num_users)
            temp_workers_loss_test = np.empty(args.num_users)
            temp_workers_loss_differenc = np.empty(args.num_users)
            temp_workers_accuracy_differenc = np.empty(args.num_users)
            flag = np.zeros(args.num_users)

            list_of_random_workers_newfl = []
            if iter < (args.epochs/2):
                for key, value in dict_workers_index.items():
                    # print(value)
                    if key == iter:
                        list_of_random_workers_newfl = dict_workers_index[key]
            else:
                list_of_random_workers_newfl = random.sample(workers, m)


            for idx in list_of_random_workers_newfl:
                #print("train started")

                # before starting train
                temp_before = copy.deepcopy(net_glob)
                # temp_before.load_state_dict(w)
                temp_before.eval()
                acc_test_local_before, loss_test_local_before = test_img(temp_before, valid_ds, args)

                local = LocalUpdate(args=args, dataset=dataset_train, idxs=dict_users[idx])
                w, loss = local.train(net=copy.deepcopy(net_glob).to(args.device))
                #print(w)
                #print("train completed")

                #print("type of idx is ", type(temp_w_locals))
                temp_w_locals.append(copy.deepcopy(w))
                temp_workers_loss[idx] = copy.deepcopy(loss)

                temp_after = copy.deepcopy(net_glob)

                temp_after.load_state_dict(w)
                temp_after.eval()
                acc_test_local_after, loss_test_local_after = test_img(temp_after, valid_ds, args)
                loss_workers_total[idx, iter] = loss_test_local_after
                temp_workers_accuracy[idx] = acc_test_local_after
                temp_workers_loss_test[idx] = loss_test_local_after
                temp_workers_loss_differenc[idx] = loss_test_local_before - loss_test_local_after
                temp_workers_accuracy_differenc[idx] = acc_test_local_after - acc_test_local_before

            print("train finished")
            while len(w_locals) < 1:
                #print("recieving started")
                index = 0
                for idx in list_of_random_workers_newfl:
                    #print("acc is ", temp_workers_accuracy[idx])
                    # print(temp_workers_loss_differenc)
                    if workers_count >= m:
                        break
                    elif temp_workers_loss_differenc[idx] >= (threshold) \
                            and temp_workers_loss_differenc[idx] > 0 \
                            and flag[idx]==0:
                        print("Update Received")
                        w_locals.append(copy.deepcopy(temp_w_locals[index]))
                        #print(temp_w_locals[index])
                        loss_locals.append(temp_workers_loss[idx])
                        flag[idx] = 1
                        workers_count += 1
                        workers_participation[idx][iter] = 1

                        data_DCFL['round_{}'.format(iter)].append({
                            'C': args.frac,
                            'User ID': idx,
                            'Loss Train': loss_workers_total[idx, iter],
                            'Loss Test': temp_workers_loss[idx],
                            'Accuracy': temp_workers_accuracy[idx]
                        })
                    index += 1
                if len(w_locals) < 1:
                    threshold = threshold / 2
                    if threshold == -np.inf:
                        threshold = 1
                print("threshold increased to ", threshold)




            # update global weights
            w_glob = FedAvg(w_locals)

            # copy weight to net_glob
            net_glob.load_state_dict(w_glob)

            #print("round completed")
            loss_avg = sum(loss_locals) / len(loss_locals)
            loss_train.append(loss_avg)
            workers_percent_dist.append(workers_count/args.num_users)


            counter_threshold_decrease[iter] = counter_threshold
            print(iter, " round dist fl finished")


            acc_test_final, loss_test_final = test_img(net_glob, dataset_test, args)
            while_counter = loss_test_final


            data_Global_DCFL["Round"].append(iter)
            data_Global_DCFL["C"].append(args.frac)
            data_Global_DCFL["Average Loss Train"].append(loss_avg)
            data_Global_DCFL["Accuracy Test"].append(Global_Accuracy_Tracker[iter])
            data_Global_DCFL["Average Loss Test"].append(Global_Loss_Tracker[iter])
            data_Global_DCFL["Workers Number"].append(workers_count)
            data_Global_DCFL["Large Test Loss"].append(float(loss_test_final))
            data_Global_DCFL["Large Test Accuracy"].append(float(acc_test_final))

            total_rounds_dcfl = iter

            iter = iter + 1


        #plot workers percent of participating
        workers_percent_final = np.zeros(args.num_users)
        workers_name = np.empty(args.num_users)
        #print(workers_participation)
        for i in range(len(workers_participation[:, 1])):
            workers_percent_final[i] = sum(workers_participation[i, :])/args.epochs
            workers_name[i] = i



        workers_fraction_each_round_newFL = sum(workers_percent_final)/len(workers_percent_final)


        # testing
        #print("testing started")
        net_glob.eval()
        #print("train test started")
        acc_train_final, loss_train_final = test_img(net_glob, dataset_train, args)
        #print("train test finished")
        acc_test_final, loss_test_final = test_img(net_glob, dataset_test, args)

        acc_full_distributed.append(acc_test_final)
        loss_full_ditributed.append(loss_test_final)

        Final_LargeDataSetTest_DCFL["C"].append(args.frac)
        Final_LargeDataSetTest_DCFL["Test Loss"].append(float(loss_test_final))
        Final_LargeDataSetTest_DCFL["Test Accuracy"].append(float(acc_test_final))
        Final_LargeDataSetTest_DCFL["Train Loss"].append(float(loss_train_final))
        Final_LargeDataSetTest_DCFL["Train Accuracy"].append(float(acc_train_final))
        Final_LargeDataSetTest_DCFL["Total Rounds"].append(int(total_rounds_dcfl))

        variable_start = variable_start + while_counter

        print("C is ", c_counter/10)

    with open('CIFAR_100users_data_main_1229-2020.json', 'w') as outfile:
        json.dump(data_main, outfile)

    with open('CIFAR_100users_data_DCFL_1229-2020.json', 'w') as outfile:
        json.dump(data_DCFL, outfile)

    with open('CIFAR_100users_data_DCFL_Global_1229-2020.json', 'w') as outfile:
        json.dump(data_Global_DCFL, outfile)

    with open('CIFAR_100users_data_main_Global_1229-2020.json', 'w') as outfile:
        json.dump(data_Global_main, outfile)

    with open('Final-CIFAR_100users_data_main_Global_1229-2020.json', 'w') as outfile:
        json.dump(Final_LargeDataSetTest_MainFL, outfile)

    with open('Final-CIFAR_100users_data_DCFL_Global_1229-2020.json', 'w') as outfile:
        json.dump(Final_LargeDataSetTest_DCFL, outfile)


    return 1