main_fair.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6

import copy
import os
import itertools
import numpy as np
from scipy.stats import mode
from torchvision import datasets, transforms, models
import torch
from torch import nn
import torch.optim as optim
from utils.sampling import fair_iid, fair_noniid
from utils.options import args_parser
from models.Update import LocalUpdate
from models.Nets import MLP, CNNMnist, CNNCifar, ResnetCifar
from models.Fed import FedAvg
from models.test import test_img, test_img_local

import pandas as pd

from sklearn import metrics
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.utils.class_weight import compute_class_weight
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader

from helpers import load_ICU_data, plot_distributions, _performance_text

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="3"

import pdb

def run_all(clf_all1, clf_all2, adv_all1, adv_all2, adv_all3):
    # 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 ICU dataset and split users
    # load ICU data set
    X, y, Z = load_ICU_data('../fairness-in-ml/data/adult.data')

    if not args.iid:
        X = X[:30000]
        y = y[:30000]
        Z = Z[:30000]

    n_points = X.shape[0]
    n_features = X.shape[1]
    n_sensitive = Z.shape[1]

    # split into train/test set
    (X_train, X_test, y_train, y_test, Z_train, Z_test) = train_test_split(X, y, Z, test_size=0.5, stratify=y, random_state=7)

    # standardize the data
    scaler = StandardScaler().fit(X_train)
    scale_df = lambda df, scaler: pd.DataFrame(scaler.transform(df), columns=df.columns, index=df.index)
    X_train = X_train.pipe(scale_df, scaler) 
    X_test = X_test.pipe(scale_df, scaler)

    class PandasDataSet(TensorDataset):
        def __init__(self, *dataframes):
            tensors = (self._df_to_tensor(df) for df in dataframes)
            super(PandasDataSet, self).__init__(*tensors)

        def _df_to_tensor(self, df):
            if isinstance(df, pd.Series):
                df = df.to_frame('dummy')
            return torch.from_numpy(df.values).float()

    def _df_to_tensor(df):
        if isinstance(df, pd.Series):
            df = df.to_frame('dummy')
        return torch.from_numpy(df.values).float()

    train_data = PandasDataSet(X_train, y_train, Z_train)
    test_data = PandasDataSet(X_test, y_test, Z_test)

    print('# train samples:', len(train_data))      # 15470
    print('# test samples:', len(test_data))

    batch_size = 32

    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, drop_last=True)
    test_loader = DataLoader(test_data, batch_size=len(test_data), shuffle=True, drop_last=True)

    # sample users
    if args.iid:
        dict_users_train = fair_iid(train_data, args.num_users)
        dict_users_test = fair_iid(test_data, args.num_users)
    else:
        train_data = [_df_to_tensor(X_train), _df_to_tensor(y_train), _df_to_tensor(Z_train)]
        test_data = [_df_to_tensor(X_test), _df_to_tensor(y_test), _df_to_tensor(Z_test)]
        #import pdb; pdb.set_trace()
        dict_users_train, rand_set_all = fair_noniid(train_data, args.num_users, num_shards=100, num_imgs=150, train=True)
        dict_users_test, _ = fair_noniid(test_data, args.num_users, num_shards=100, num_imgs=150, train=False, rand_set_all=rand_set_all)

    train_data = [_df_to_tensor(X_train), _df_to_tensor(y_train), _df_to_tensor(Z_train)]
    test_data = [_df_to_tensor(X_test), _df_to_tensor(y_test), _df_to_tensor(Z_test)]

    class LocalClassifier(nn.Module):
        def __init__(self, n_features, n_hidden=32, p_dropout=0.2):
            super(LocalClassifier, self).__init__()
            self.network1 = nn.Sequential(
                nn.Linear(n_features, n_hidden),
                nn.ReLU(),
                nn.Dropout(p_dropout),
                nn.Linear(n_hidden, n_hidden),
                nn.ReLU(),
                nn.Dropout(p_dropout),
                nn.Linear(n_hidden, n_hidden)
            )
            self.network2 = nn.Sequential(
                nn.ReLU(),
                nn.Dropout(p_dropout),
                nn.Linear(n_hidden, 1)
            )

        def forward(self, x):
            mid = self.network1(x)
            final = torch.sigmoid(self.network2(mid))
            return mid, final

    def pretrain_classifier(clf, data_loader, optimizer, criterion):
        losses = 0.0
        for x, y, _ in data_loader:
            x = x.to(args.device)
            y = y.to(args.device)
            clf.zero_grad()
            mid, p_y = clf(x)
            loss = criterion(p_y, y)
            loss.backward()
            optimizer.step()
            losses += loss.item()
        print ('loss', losses/len(data_loader))
        return clf

    def test_classifier(clf, data_loader):
        losses = 0
        assert len(data_loader) == 1
        with torch.no_grad():
            for x, y_test, _ in data_loader:
                x = x.to(args.device)
                mid, y_pred = clf(x)
                y_pred = y_pred.cpu()
                clf_accuracy = metrics.accuracy_score(y_test, y_pred > 0.5) * 100
        return clf_accuracy

    class Adversary(nn.Module):

        def __init__(self, n_sensitive, n_hidden=32):
            super(Adversary, self).__init__()
            self.network = nn.Sequential(
                nn.Linear(n_hidden, n_hidden),
                nn.ReLU(),
                nn.Linear(n_hidden, n_hidden),
                nn.ReLU(),
                nn.Linear(n_hidden, n_hidden),
                nn.ReLU(),
                nn.Linear(n_hidden, n_sensitive),
            )

        def forward(self, x):
            return torch.sigmoid(self.network(x))

    def pretrain_adversary(adv, clf, data_loader, optimizer, criterion):
        losses = 0.0
        for x, _, z in data_loader:
            x = x.to(args.device)
            z = z.to(args.device)
            mid, p_y = clf(x)
            mid = mid.detach()
            p_y = p_y.detach()
            adv.zero_grad()
            p_z = adv(mid)
            loss = (criterion(p_z.to(args.device), z.to(args.device)) * lambdas.to(args.device)).mean()
            loss.backward()
            optimizer.step()
            losses += loss.item()
        print ('loss', losses/len(data_loader))
        return adv

    def test_adversary(adv, clf, data_loader):
        losses = 0
        adv_accuracies = []
        assert len(data_loader) == 1
        with torch.no_grad():
            for x, _, z_test in data_loader:
                x = x.to(args.device)
                mid, p_y = clf(x)
                mid = mid.detach()
                p_y = p_y.detach()
                p_z = adv(mid)
                for i in range(p_z.shape[1]):
                    z_test_i = z_test[:,i]
                    z_pred_i = p_z[:,i]
                    z_pred_i = z_pred_i.cpu()
                    adv_accuracy = metrics.accuracy_score(z_test_i, z_pred_i > 0.5) * 100
                    adv_accuracies.append(adv_accuracy)
        return adv_accuracies

    def train_both(clf, adv, data_loader, clf_criterion, adv_criterion, clf_optimizer, adv_optimizer, lambdas):
        # Train adversary
        adv_losses = 0.0
        for x, y, z in data_loader:
            x = x.to(args.device)
            z = z.to(args.device)
            local, p_y = clf(x)
            adv.zero_grad()
            p_z = adv(local)
            loss_adv = (adv_criterion(p_z.to(args.device), z.to(args.device)) * lambdas.to(args.device)).mean()
            loss_adv.backward()
            adv_optimizer.step()
            adv_losses += loss_adv.item()
        print ('adversarial loss', adv_losses/len(data_loader))

        # Train classifier on single batch
        clf_losses = 0.0
        for x, y, z in data_loader:
            pass
        x = x.to(args.device)
        y = y.to(args.device)
        z = z.to(args.device)
        local, p_y = clf(x)
        p_z = adv(local)
        clf.zero_grad()
        if args.adv:
            clf_loss = clf_criterion(p_y.to(args.device), y.to(args.device)) - (adv_criterion(p_z.to(args.device), z.to(args.device)) * lambdas.to(args.device)).mean()
        else:
            clf_loss = clf_criterion(p_y.to(args.device), y.to(args.device))
        clf_loss.backward()
        clf_optimizer.step()
        clf_losses += clf_loss.item()
        print ('classifier loss', clf_losses/len(data_loader))
        return clf, adv

    def eval_performance_text(test_loader_i, local_clf_i, adv_i):
        with torch.no_grad():
            for test_x, test_y, test_z in test_loader_i:
                test_x = test_x.to(args.device)
                local_pred, clf_pred = local_clf_i(test_x)
                adv_pred = adv_i(local_pred)

            y_post_clf = pd.Series(clf_pred.cpu().numpy().ravel(), index=y_test[list(dict_users_train[idx])].index)
            Z_post_adv = pd.DataFrame(adv_pred.cpu().numpy(), columns=Z_test.columns)

            clf_roc_auc,clf_accuracy,adv_acc1,adv_acc2,adv_roc_auc = _performance_text(test_y, test_z, y_post_clf, Z_post_adv, epoch=None)
        return clf_roc_auc,clf_accuracy,adv_acc1,adv_acc2,adv_roc_auc

    def eval_global_performance_text(test_loader_i, local_clf_i, adv_i, global_clf):
        with torch.no_grad():
            for test_x, test_y, test_z in test_loader_i:
                test_x = test_x.to(args.device)
                local_pred, clf_pred = local_clf_i(test_x)
                adv_pred = adv_i(local_pred)
                global_pred = global_clf(local_pred)

            y_post_clf = pd.Series(global_pred.cpu().numpy().ravel(), index=y_test[list(dict_users_train[idx])].index)
            Z_post_adv = pd.DataFrame(adv_pred.cpu().numpy(), columns=Z_test.columns)

            clf_roc_auc,clf_accuracy,adv_acc1,adv_acc2,adv_roc_auc = _performance_text(test_y, test_z, y_post_clf, Z_post_adv, epoch=None)
        return clf_roc_auc,clf_accuracy,adv_acc1,adv_acc2,adv_roc_auc

    lambdas = torch.Tensor([30.0, 30.0])
    net_local_list = []

    print ('\n\n======================== STARTING LOCAL TRAINING ========================\n\n\n')

    for idx in range(args.num_users):
        print ('\n======================== LOCAL TRAINING, USER %d ========================\n\n\n' %idx)
        train_data_i_raw = [torch.FloatTensor(bb[list(dict_users_train[idx])]) for bb in train_data]
        train_data_i = TensorDataset(train_data_i_raw[0],train_data_i_raw[1],train_data_i_raw[2])
        train_loader_i = torch.utils.data.DataLoader(train_data_i, batch_size=batch_size, shuffle=False, num_workers=4)

        test_data_i_raw = [torch.FloatTensor(bb[list(dict_users_train[idx])]) for bb in test_data]
        test_data_i = TensorDataset(test_data_i_raw[0],test_data_i_raw[1],test_data_i_raw[2])
        test_loader_i = torch.utils.data.DataLoader(test_data_i, batch_size=len(test_data_i), shuffle=False, num_workers=4)

        local_clf_i = LocalClassifier(n_features=n_features).to(args.device)
        local_clf_criterion_i = nn.BCELoss().to(args.device)
        local_clf_optimizer_i = optim.SGD(local_clf_i.parameters(), lr=0.1)

        adv_i = Adversary(Z_train.shape[1]).to(args.device)
        adv_criterion_i = nn.BCELoss(reduce=False).to(args.device)
        adv_optimizer_i = optim.SGD(adv_i.parameters(), lr=0.1)

        net_local_list.append([train_loader_i,test_loader_i,local_clf_i,local_clf_optimizer_i,local_clf_criterion_i,adv_i,adv_criterion_i,adv_optimizer_i])
        
        N_CLF_EPOCHS = 10
        for epoch in range(N_CLF_EPOCHS):
            print ('======================== pretrain_classifier epoch %d ========================' %epoch)
            local_clf = pretrain_classifier(local_clf_i, train_loader_i, local_clf_optimizer_i, local_clf_criterion_i)
        # test classifier
        # print ('\npretrained test accuracy on income prediction', test_classifier(local_clf_i, test_loader))
        # print ()
        print ('======================== local classifier pretraining: evaluating _performance_text on device %d ========================' %idx)
        eval_performance_text(test_loader_i, local_clf_i, adv_i)

        N_ADV_EPOCHS = 10
        for epoch in range(N_ADV_EPOCHS):
            print ('======================== pretrain_adversary epoch %d ========================' %epoch)
            pretrain_adversary(adv_i, local_clf_i, train_loader_i, adv_optimizer_i, adv_criterion_i)

        # test adversary
        # print ('\npretrained adversary accuracy on race, sex prediction', test_adversary(adv_i, local_clf_i, test_loader))
        # print ()
        print ('======================== local adversary pretraining: evaluating _performance_text on device %d ========================' %idx)
        eval_performance_text(test_loader_i, local_clf_i, adv_i)
        print ('======================== by now both the local classifier and the local adversary should do well ========================')

        # train both
        N_EPOCH_COMBINED = 0 #250
        for epoch in range(N_EPOCH_COMBINED):
            print ('======================== combined training epoch %d ========================' %epoch)
            clf, adv = train_both(local_clf_i, adv_i, train_loader_i, local_clf_criterion_i, adv_criterion_i,
                                  local_clf_optimizer_i, adv_optimizer_i, lambdas)
            # test classifier
            #print ('final test accuracy on income prediction', test_classifier(clf, test_loader))
            # test adversary
            #print ('final adversary accuracy on race, sex prediction', test_adversary(adv, clf, test_loader))

        print ('======================== local classifier and adversary pretraining: evaluating _performance_text on device %d ========================' %idx)
        eval_performance_text(test_loader_i, local_clf_i, adv_i)

        print ('======================== by now the local classifier should do well but the local adversary should not do well ========================')

    print ('======================== done pretraining local classifiers and adversaries ========================')

    class GlobalClassifier(nn.Module):
        def __init__(self, n_hidden=32, p_dropout=0.2):
            super(GlobalClassifier, self).__init__()
            self.global_network = nn.Sequential(
                nn.Linear(n_hidden, n_hidden),
                nn.ReLU(),
                nn.Dropout(p_dropout),
                nn.Linear(n_hidden, n_hidden),
                nn.ReLU(),
                nn.Dropout(p_dropout),
                nn.Linear(n_hidden, 1),
            )

        def forward(self, local):
            final = torch.sigmoid(self.global_network(local))
            return final

    # build global model
    global_clf = GlobalClassifier().to(args.device)
    global_clf_criterion = nn.BCELoss().to(args.device)
    global_clf_optimizer = optim.Adam(global_clf.parameters(), lr=0.01)

    # copy weights
    w_glob = global_clf.state_dict()

    print ('\n\n======================== STARTING GLOBAL TRAINING ========================\n\n\n')

    global_epochs = 10
    for iter in range(global_epochs):
        w_locals, loss_locals = [], []
        for idx in range(args.num_users):
            print ('\n\n======================== GLOBAL TRAINING, ITERATION %d, USER %d ========================\n\n\n' %(iter,idx))
            train_loader_i,test_loader_i,local_clf_i,local_clf_optimizer_i,local_clf_criterion_i,adv_i,adv_criterion_i,adv_optimizer_i = net_local_list[idx]
            # train both local models: classifier and adversary
            if iter % 2 == 0:
                N_EPOCH_COMBINED = 0 #65
                for epoch in range(N_EPOCH_COMBINED):
                    print ('======================== combined training epoch %d ========================' %epoch)
                    local_clf_i, adv_i = train_both(local_clf_i, adv_i, train_loader_i, local_clf_criterion_i, adv_criterion_i,
                                                    local_clf_optimizer_i, adv_optimizer_i, lambdas)

            local = LocalUpdate(args=args, dataset=train_loader_i)
            w, loss = local.train(local_net=local_clf_i, local_opt=local_clf_optimizer_i, local_adv=adv_i, adv_opt=adv_optimizer_i, global_net=copy.deepcopy(global_clf).to(args.device), global_opt=global_clf_optimizer)

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

        w_glob = FedAvg(w_locals)
        # copy weight to net_glob
        global_clf.load_state_dict(w_glob)

    
    for idx in range(args.num_users):
        train_loader_i,test_loader_i,local_clf_i,local_clf_optimizer_i,local_clf_criterion_i,adv_i,adv_criterion_i,adv_optimizer_i = net_local_list[idx]

        print ('======================== local and global training: evaluating _performance_text on device %d ========================' %idx)
        eval_performance_text(test_loader_i, local_clf_i, adv_i)
        print ('======================== by now the local classifier should do well but the local adversary should not do well ========================')

        print ('======================== local and global training: evaluating _global_performance_text on device %d ========================' %idx)
        clf_roc_auc,clf_accuracy,adv_acc1,adv_acc2,adv_roc_auc = eval_global_performance_text(test_loader_i, local_clf_i, adv_i, global_clf)
        print ('======================== by now the global classifier should work better than local classifier ========================')

        clf_all1.append(clf_roc_auc)
        clf_all2.append(clf_accuracy)
        adv_all1.append(adv_acc1)
        adv_all2.append(adv_acc2)
        adv_all3.append(adv_roc_auc)
    print ('clf_all1', np.mean(np.array(clf_all1)), np.std(np.array(clf_all1)))
    print ('clf_all2', np.mean(np.array(clf_all2)), np.std(np.array(clf_all2)))
    print ('adv_all1', np.mean(np.array(adv_all1)), np.std(np.array(adv_all1)))
    print ('adv_all2', np.mean(np.array(adv_all2)), np.std(np.array(adv_all2)))
    print ('adv_all3', np.mean(np.array(adv_all3)), np.std(np.array(adv_all3)))
    return clf_all1, clf_all2, adv_all1, adv_all2, adv_all3


if __name__ == '__main__':
    clf_all1, clf_all2, adv_all1, adv_all2, adv_all3 = [], [], [], [], []
    for _ in range(10):
        clf_all1, clf_all2, adv_all1, adv_all2, adv_all3 = run_all(clf_all1, clf_all2, adv_all1, adv_all2, adv_all3)
    print ('final')
    print ('clf_all1', np.mean(np.array(clf_all1)), np.std(np.array(clf_all1)))
    print ('clf_all2', np.mean(np.array(clf_all2)), np.std(np.array(clf_all2)))
    print ('adv_all1', np.mean(np.array(adv_all1)), np.std(np.array(adv_all1)))
    print ('adv_all2', np.mean(np.array(adv_all2)), np.std(np.array(adv_all2)))
    print ('adv_all3', np.mean(np.array(adv_all3)), np.std(np.array(adv_all3)))