main_targeted_attack.py

import numpy as np
import json
import torch
import torch.optim as optim
import torch.nn as nn
import argparse
import logging
import os
import copy
import datetime
import random

from functools import reduce
from model import *
from utils import *
from defense_utils import *


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, default='cifar10', help='dataset used for training')
    parser.add_argument('--net_config', type=lambda x: list(map(int, x.split(', '))))
    parser.add_argument('--partition', type=str, default='noniid', help='the data partitioning strategy')
    parser.add_argument('--batch-size', type=int, default=64, help='input batch size for training')
    parser.add_argument('--lr', type=float, default=0.01, help='learning rate (default: 0.1)')
    parser.add_argument('--epochs', type=int, default=1, help='number of local epochs')
    parser.add_argument('--n_parties', type=int, default=20, help='number of workers in a distributed cluster')
    parser.add_argument('--comm_round', type=int, default=100, help='number of maximum communication roun')
    parser.add_argument('--init_seed', type=int, default=0, help="Random seed")
    parser.add_argument('--datadir', type=str, required=False, default="./data/", help="Data directory")
    parser.add_argument('--reg', type=float, default=1e-5, help="L2 regularization strength")
    parser.add_argument('--logdir', type=str, required=False, default="./logs/", help='Log directory path')
    parser.add_argument('--modeldir', type=str, required=False, default="./models/", help='Model directory path')
    parser.add_argument('--beta', type=float, default=0.5, help='The parameter for the dirichlet distribution for data partitioning')
    parser.add_argument('--device', type=str, default='cuda:0', help='The device to run the program')
    parser.add_argument('--log_file_name', type=str, default=None, help='The log file name')
    parser.add_argument('--optimizer', type=str, default='sgd', help='the optimizer')    
    parser.add_argument('--global_defense', type=str, default='average',help='communication strategy')
    
    # poison settings
    parser.add_argument('--attacker_ratio', type=float, default=0.2, help='ratio for number of attackers')
    parser.add_argument('--poisoning_rate', type=float, default=0.8, help='poisoning portion')
    parser.add_argument('--trigger_label', type=int, default=1, help='The NO. of trigger label')
    parser.add_argument('--trigger_path', default="./triggers/trigger_white.png", help='Trigger Path')
    parser.add_argument('--trigger_size', type=int, default=5, help='Trigger Size')
        
    return parser.parse_args()


def init_nets(net_configs, n_parties, args, device='cpu'):
    nets = {net_i: None for net_i in range(n_parties)}
    if args.dataset in {'cifar10', 'svhn'}:
        n_classes = 10
    elif args.dataset == 'tinyimagenet':
        n_classes = 200

    for net_i in range(n_parties):
        if args.dataset in {'cifar10', 'svhn'}:
            net = ResNet18_cifar10()
        elif args.dataset == 'tinyimagenet':
            net = ResNet18_tinyimage()

        if device == 'cpu':
            net.to(device)
        else:
            net = net.cuda()
        nets[net_i] = net


    model_meta_data = []
    layer_type = []
    for (k, v) in nets[0].state_dict().items():
        model_meta_data.append(v.shape)
        layer_type.append(k)

    return nets, model_meta_data, layer_type

def train_net(net_id, net, train_dataloader, test_dataloader, epochs, lr, args_optimizer, args, device="cpu"):
    net.cuda()
    if args_optimizer == 'adam':
        optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=lr, weight_decay=args.reg)
    elif args_optimizer == 'amsgrad':
        optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=lr, weight_decay=args.reg,
                               amsgrad=True)
    elif args_optimizer == 'sgd':
        optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()), lr=lr, momentum=0.9,
                              weight_decay=args.reg)
    criterion = nn.CrossEntropyLoss().cuda()

    cnt = 0
    for epoch in range(epochs):
        epoch_loss_collector = []
        for batch_idx, (x, target) in enumerate(train_dataloader):
            x, target = x.cuda(), target.cuda()

            optimizer.zero_grad()
            x.requires_grad = False
            target.requires_grad = False
            target = target.long()

            out = net(x)
            loss = criterion(out, target)

            loss.backward()
            optimizer.step()

            cnt += 1
            epoch_loss_collector.append(loss.item())

        epoch_loss = sum(epoch_loss_collector) / len(epoch_loss_collector)
        
    train_acc, _ = compute_accuracy(net, train_dataloader, device=device)
    test_acc, conf_matrix, _ = compute_accuracy(net, test_dataloader, get_confusion_matrix=True, device=device)
    net.to('cpu')

    return train_acc, test_acc


def local_train_net(nets, args, net_dataidx_map, attacker_id_list=[], train_dl=None, 
                    test_dl=None, global_model = None, round=None, device="cpu"):
    avg_acc = 0.0
    acc_list = []
    if global_model:
        global_model.cuda()

    for net_id, net in nets.items():
        dataidxs = net_dataidx_map[net_id]
        if net_id in attacker_id_list:
            prefix = 'attacker'
            train_dl_local, test_dl_local, _, _ = get_poisoned_dataloader(args, args.dataset, args.datadir, args.batch_size, 32, dataidxs)
        else:
            prefix = 'normal'
            train_dl_local, test_dl_local, _, _ = get_dataloader(args.dataset, args.datadir, args.batch_size, 32, dataidxs)
        train_dl_global, test_dl_global, _, _ = get_dataloader(args.dataset, args.datadir, args.batch_size, 32)
        n_epoch = args.epochs

        if net_id in attacker_id_list:
            trainacc, testacc = train_net(net_id, net, train_dl_local, test_dl, n_epoch, args.lr, args.optimizer, args, device=device)
        else:
            trainacc, testacc = train_net(net_id, net, train_dl_local, test_dl, n_epoch, args.lr, args.optimizer, args, device=device)
            
        logger.info(prefix + " net %d final test acc %f" % (net_id, testacc))
        avg_acc += testacc
        acc_list.append(testacc)
    avg_acc /= args.n_parties
    
    if global_model:
        global_model.to('cpu')
    return nets


args = get_args()
mkdirs(args.logdir)
mkdirs(args.modeldir)

device = torch.device(args.device)
for handler in logging.root.handlers[:]:
    logging.root.removeHandler(handler)

if args.log_file_name is None:
    args.log_file_name = f'{args.dataset}-beta_{args.beta}-num-party_{args.n_parties}-local-epoch_{args.epochs}-type_targeted-attacker-ratio_{args.attacker_ratio}-{args.global_defense}-poisoning_{args.poisoning_rate}'
log_path = args.log_file_name + '.log'
logging.basicConfig(
    filename=os.path.join(args.logdir, log_path),
    format='%(asctime)s %(levelname)-8s %(message)s',
    datefmt='%m-%d %H:%M', level=logging.DEBUG, filemode='w')

logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
logging.getLogger('PIL').setLevel(logging.WARNING)
logger.info(device)

seed = args.init_seed
logger.info("#" * 100)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed(seed)
random.seed(seed)

logger.info("Partitioning data")
X_train, y_train, X_test, y_test, net_dataidx_map, traindata_cls_counts = partition_data(
    args.dataset, args.datadir, args.logdir, args.partition, args.n_parties, beta=args.beta)

n_party_per_round = min(args.n_parties, 10)
party_list = [i for i in range(args.n_parties)]
party_list_rounds = []
if n_party_per_round != args.n_parties:
    for i in range(args.comm_round):
        this_round = random.sample(party_list, n_party_per_round)
        this_round.sort()
        party_list_rounds.append(this_round)
else:
    for i in range(args.comm_round):
        party_list_rounds.append(party_list)


n_attacker = int(args.attacker_ratio * args.n_parties)
expected_n_attacker = int(args.attacker_ratio * n_party_per_round)
args.expected_n_attacker = expected_n_attacker - 1
attacker_id_list = random.sample(party_list, n_attacker)      
logger.info(">> Attacker Network IDX: {}".format(' '.join(map(str, attacker_id_list))))

    
n_classes = len(np.unique(y_train))

train_dl_global, test_dl, train_ds_global, test_ds_global = get_dataloader(args.dataset,
                                                                           args.datadir,
                                                                           args.batch_size,
                                                                           32)

train_dl_global_poisoned, test_dl_poisoned, train_ds_global_poisoned, test_ds_global_poisoned = get_poisoned_dataloader(args, args.dataset,
                                                                           args.datadir,
                                                                           args.batch_size,
                                                                           32)

print("len train_dl_global:", len(train_ds_global))
train_dl=None
data_size = len(test_ds_global)

logger.info("Initializing nets")
nets, local_model_meta_data, layer_type = init_nets(args.net_config, args.n_parties, args, device='cpu')

global_models, global_model_meta_data, global_layer_type = init_nets(args.net_config, 1, args, device='cpu')
global_model = global_models[0]
n_comm_rounds = args.comm_round         
    
for round in range(n_comm_rounds):
    logger.info("in comm round:" + str(round))
    party_list_this_round = party_list_rounds[round]

    global_w = global_model.state_dict()
    
    
    if round != 0:
        prev_prev_global_w = copy.deepcopy(prev_global_w)
        prev_global_w = copy.deepcopy(global_model.state_dict())
    
        
    nets_this_round = {k: nets[k] for k in party_list_this_round}
    for net in nets_this_round.values():
        net.load_state_dict(global_w)
        
    local_train_net(nets_this_round, args, net_dataidx_map, attacker_id_list=attacker_id_list, 
                    train_dl=train_dl, test_dl=test_dl, global_model= global_model, device=device)

    total_data_points = sum([len(net_dataidx_map[r]) for r in party_list_this_round])
    fed_avg_freqs = [len(net_dataidx_map[r]) / total_data_points for r in party_list_this_round]
    
    if round == 0:
        prev_prev_global_w = copy.deepcopy(global_model.state_dict())
        prev_global_w = copy.deepcopy(global_model.state_dict())
        
        for net_id, net in enumerate(nets_this_round.values()):
            net_para = net.state_dict()
            if net_id == 0:
                for key in net_para:
                    prev_global_w[key] = net_para[key] / len(nets_this_round)
            else:
                for key in net_para:
                    prev_global_w[key] += net_para[key] / len(nets_this_round)
    
    global_w = global_aggregation(nets_this_round, args, fed_avg_freqs, global_w, party_list_this_round, prev_global_w, prev_prev_global_w, global_model, round, logger)
    global_model.load_state_dict(global_w)
    global_model.cuda()
    
    test_acc, conf_matrix, _ = compute_accuracy(global_model, test_dl, get_confusion_matrix=True, device=device)
    logger.info('>> Global Model Test accuracy: %f' % test_acc)
    
    asr_acc,  _ = compute_accuracy(global_model, test_dl_poisoned, device=device)
    logger.info('>> Global Model Attack Success Rate: %f' % asr_acc)
        
    global_model.to('cpu')
    
mkdirs(args.modeldir+'fedavg/')
torch.save(global_model.state_dict(), args.modeldir+'fedavg/'+'globalmodel'+args.log_file_name+'.pth')
torch.save(nets[0].state_dict(), args.modeldir+'fedavg/'+'localmodel0'+args.log_file_name+'.pth')