fed_dual_att.py

# -*- coding: utf-8 -*-
"""
-----------------------------------------------
# File: fed_dual_att.py
# This file is created by Chuanting Zhang
# Email: chuanting.zhang@kaust.edu.sa
# Date: 2020-07-25 (YYYY-MM-DD)
-----------------------------------------------
"""
import numpy as np
import h5py
import tqdm
import copy
import torch
import pandas as pd
import random
from collections import defaultdict
from torch.utils.data import DataLoader
import os
import sys
from sklearn import metrics
from scipy.spatial.distance import pdist

sys.path.append('../')
from utils.misc import args_parser, avg_dual_att
from utils.misc import get_data, process_isolated, get_warm_up_data, set_logger
from utils.misc import get_cluster_label, jfi, cv
from utils.models import LSTM
from utils.fed_update import LocalUpdate, test_inference

from utils.attacks import ModelPoison, AdaptiveModelPoison
from utils.defenses import RobustAggregation

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


def get_stat_mean(d):
    d = d.iloc[:-24 * args.test_days, :]
    df_avg = d.groupby([d.index.week, d.index.hour]).mean().reset_index().iloc[:, 2:]
    df_avg = (df_avg - df_avg.mean()) / df_avg.std()
    return copy.deepcopy(df_avg.T)


def get_warm_model(g_index, save_name):
    # warm-up model using the statistical mean model
    warm_xc, warm_xp, warm_y = [], [], []
    # print(g_index)

    model = LSTM(args).to(device)
    for i in g_index:
        cell_xc, cell_xp, cell_y = get_warm_up_data(args, df.loc[i][:-1])
        if args.phi > 0:
            n_transfer = int(np.floor(len(cell_xc) * args.phi))
            idx = [a for a in np.random.randint(0, len(cell_xc), n_transfer)]
            warm_xc.append(cell_xc[idx])
            warm_xp.append(cell_xp[idx])
            warm_y.append(cell_y[idx])
        else:
            warm_xc.append(cell_xc)
            warm_xp.append(cell_xp)
            warm_y.append(cell_y)
    warm_xc_arr = np.concatenate(warm_xc, axis=0)[:, :, np.newaxis]
    if args.period_size > 0:
        warm_xp_arr = np.concatenate(warm_xp, axis=0)[:, :, np.newaxis]
    else:
        warm_xp_arr = warm_xc_arr
    warm_y_arr = np.concatenate(warm_y, axis=0)

    warm_data = list(zip(*[warm_xc_arr, warm_xp_arr, warm_y_arr]))
    warm_loader = DataLoader(warm_data, shuffle=False, batch_size=args.batch_size)
    warm_criterion = torch.nn.MSELoss().to(device)
    if args.opt == 'adam':
        warm_opt = torch.optim.Adam(model.parameters(), lr=args.w_lr)
    elif args.opt == 'sgd':
        warm_opt = torch.optim.SGD(model.parameters(), lr=args.w_lr, momentum=args.momentum)

    warm_scheduler = torch.optim.lr_scheduler.MultiStepLR(warm_opt, milestones=[0.5 * args.w_epoch,
                                                                                0.75 * args.w_epoch],
                                                          gamma=0.1)

    for epoch in range(args.w_epoch):
        warm_epoch_loss = []
        model.train()
        for batch_idx, (xc, xp, y) in enumerate(warm_loader):
            xc, xp, y = xc.float().to(device), xp.float().to(device), y.float().to(device)
            model.zero_grad()
            pred = model(xc, xp)
            loss = warm_criterion(y, pred)
            warm_epoch_loss.append(loss)
            loss.backward()
            warm_opt.step()

        warm_scheduler.step()

    return model.state_dict()


if __name__ == '__main__':
    args = args_parser()

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    if not os.path.isdir(args.directory):
        os.mkdir(args.directory)

    data, df_ori, selected_cells, mean, std, lng, lat = get_data(args)
    # print(selected_cells)
    device = 'cuda' if args.gpu else 'cpu'

    parameter_list = 'FedDualAtt-data-{:}-type-{:}-'.format(args.file, args.type)
    parameter_list += 'rho-{:.3f}-cluster-{:}-lr-{:.4f}-'.format(args.rho, args.cluster, args.lr)
    parameter_list += '-frac-{:.2f}-le-{:}-lb-{:}-epsilon-{:.2f}-seed-{:}-'.format(args.frac, args.local_epoch,
                                                                                   args.local_bs, args.epsilon,
                                                                                   args.seed)
    parameter_list += 'warm_up:{:}'.format(args.warm_up)

    log_id = parameter_list
    logger = set_logger(log_id=log_id, args=args, logger_name=__name__)

    # print(args)

    train, val, test = process_isolated(args, data)
    # get the statistical mean of the traffic data
    df = get_stat_mean(data)
    # print(df_mean.head())
    data_dist = pdist(df.values)
    data_jfi = jfi(np.array(data_dist))
    data_cv = cv(np.array(data_dist))
    # print('jfi: {:.4f}, cv: {:.4f}'.format(data_jfi, data_cv))

    # dual-stage iterative clustering
    df['label'] = get_cluster_label(args, df, lng, lat)

    global_model = LSTM(args).to(device)
    # use this warm-up model as initialization
    cluster_weights = defaultdict()

    if args.warm_up:
        warm_weights = copy.deepcopy(get_warm_model(selected_cells, log_id))
        global_weights = copy.deepcopy(warm_weights)
        global_model.load_state_dict(global_weights)
        for label in df['label'].unique():
            cluster_weights[label] = copy.deepcopy(warm_weights)
    else:
        warm_weights = copy.deepcopy(global_model.state_dict())
        for label in df['label'].unique():
            cluster_weights[label] = copy.deepcopy(warm_weights)

    # training
    best_val_loss = None
    val_loss = []
    val_acc = []

    if args.poison:
        m = max(int(args.adversary_frac * args.bs), 1)
        adversary_idx = random.sample(selected_cells, m)
        historic_global_model = None
    if args.apply_defense == 'multi_krum' or args.apply_defense == 'trimmed_mean':
        defense = RobustAggregation(args)

    ## attack model availability
    target = None

    for epoch in tqdm.tqdm(range(args.epochs)):
        local_weights, local_losses = defaultdict(list), defaultdict(list)
        # print(f'\n Global Training Round: {epoch+1}|\n')

        m = max(int(args.frac * args.bs), 1)
        cell_idx = random.sample(selected_cells, m)

        historic_global_model_temp = copy.deepcopy(global_model)

        avg_loss = 0
        for cell in cell_idx:
            group_id = df.loc[cell]['label']
            # print('Group ID:', group_id)
            global_model.load_state_dict(global_weights)

            cell_train, cell_test = train[cell], test[cell]

            if args.poison and cell in adversary_idx:
                local_model = AdaptiveModelPoison(args, cell_train, cell_test,
                                                    historic_global_model, target, args.apply_defense)
            else:
                # print('normal cell: {}'.format(cell))
                local_model = LocalUpdate(args, cell_train, cell_test)

            w, loss, epoch_loss = local_model.update_weights(model=copy.deepcopy(global_model),
                                                             global_round=epoch)
            avg_loss += loss

            local_weights[group_id].append(copy.deepcopy(w))
            local_losses[group_id].append(copy.deepcopy(loss))

        # Update global model
        local_cluster = defaultdict()
        for group_id in local_weights.keys():
            local_cluster[group_id] = avg_dual_att(local_weights[group_id], cluster_weights[group_id],
                                                      warm_weights,
                                                      args.epsilon, args.rho)
        cw = []
        for c_key, c_weights in local_cluster.items():
            cw.append(c_weights)
        global_weights = avg_dual_att(cw, global_weights, warm_weights, args.epsilon, args.rho)
        # global_weights = average_weights(cw)
        global_model.load_state_dict(global_weights)

        historic_global_model = historic_global_model_temp

    pred, truth = defaultdict(), defaultdict()
    test_loss_list = []
    test_mse_list = []
    nrmse = 0.0

    global_model.load_state_dict(global_weights)

    with torch.no_grad():
        for cell in selected_cells:
            cell_test = test[cell]
            group_id = int(df.loc[cell]['label'])

            test_loss, test_mse, test_nrmse, pred[cell], truth[cell] = test_inference(args, global_model, cell_test)
            # print(f'Cell: {cell} MSE: {test_mse:.4f}')
            nrmse += test_nrmse

            test_loss_list.append(test_loss)
            test_mse_list.append(test_mse)

    df_pred = pd.DataFrame.from_dict(pred)
    df_truth = pd.DataFrame.from_dict(truth)

    mse = metrics.mean_squared_error(df_pred.values.ravel(), df_truth.values.ravel())
    mae = metrics.mean_absolute_error(df_pred.values.ravel(), df_truth.values.ravel())
    nrmse = nrmse / len(selected_cells)
    logger.info('FedDualAtt File: {:}, Type: {:}, BS: {:}, frac: {:.2f}, Cluster: {:}, rho: '
                '{:.2f}, epsilon: {:.2f}, seed: {:}, lb: {:}, le: {:}, close: {:}, period: {:}, hidden: {:}, layers: {:},'
                ' lr: {:.4f}, w_lr: {:.4f}, '
                'MSE: {:.4f}, MAE: {:.4f}, NRMSE: {:.4f}'.format(
                    args.file, args.type, args.bs, args.frac,
                    args.cluster, args.rho, args.epsilon,
                    args.seed, args.local_bs, args.local_epoch, args.close_size, args.period_size,
                    args.hidden_dim, args.phi, args.lr, args.w_lr,
                    mse, mae, nrmse))