poisoned_train.py

import copy
import json
import time

from tqdm import tqdm

from Models import Client, Server
import torch
import numpy as np

from Models.ModelFactory import get_bottom_and_top_model
from Utils import DatasetUtil
from Utils.FlexLR import FlexLearningRate
from Utils.LogUtil import LogUtil
from Utils.PrintUtil import getNowString, Logger, calculate_persample_transmission_cost, dirCheck
import os, sys

from Utils.args import parser_args


def poisoned_train(args):

    dataset_name = args.dataset

    start_time = time.time()
    now_str = getNowString()
    logUtil = LogUtil(now_str)

    num_users = args.num_clients
    epochs = args.global_epochs
    local_epochs = args.local_epochs
    batch_size = args.batch_size
    frac = args.frac
    lr = args.lr
    wm_engage = args.wm_engage
    backdoor_engage = args.bd_engage
    backdoor_type = args.bd_type
    dp_epsilon = args.epsilon
    flex_lr_stages = args.lr_stages
    flex_lr_gamma = args.gamma

    noniid_mode = args.noniid
    noniid_beta = args.beta

    grad_clip = args.grad_clip

    bd_intensity = args.bd_intensity

    rho = args.rho

    model_name = args.model
    split_point = args.split_point



    bd_overlap = True if args.bd_overlap == 1 else False



    cos_annealing_multiplier = args.cos_annealing_multiplier

    if model_name == "ResNet18":
        with open(os.path.join('wm_conf', 'resnet18_wm_cfg.json'), "r") as f:
            wm_config = json.load(f)
    elif model_name == "MobileNet":
        with open(os.path.join('wm_conf', 'mobilenet_wm_cfg.json'), "r") as f:
            wm_config = json.load(f)
    elif model_name == "DenseNet":
        with open(os.path.join('wm_conf', 'densenet_wm_cfg.json'), "r") as f:
            wm_config = json.load(f)

    consoleLogDir = logUtil.getDir()


    log_file = os.path.join(consoleLogDir, f"{dataset_name}_{now_str}_{epochs}_{wm_engage}_{backdoor_engage}.log")

    # 重定向 stdout
    sys.stdout = Logger(log_file, "a")

    train_dataset, test_dataset, num_classes = DatasetUtil.get_datasets(dataset_name=dataset_name,
                                                        use_augmentation=True)

    net_globe_client, top_model = get_bottom_and_top_model(model_name=model_name, split_point=split_point,
                                                           num_classes=num_classes)
    # top_model = None
    # if model_name == "ResNet18":
    #     # net_globe_client = SFL_ResNet18.ResNet_Bottom()
    #     model = ResNet18()
    #     net_globe_client, top_model = get_bottom_and_top_model(model_name=model_name, split_point=split_point, num_classes=num_classes)
    # elif model_name == "MobileNet":
    #     net_globe_client = MobileNetV2_Bottom()

    transmission_cost_per_sample = calculate_persample_transmission_cost(bottom_model=net_globe_client, dataset=test_dataset)
    print(f"transmission_cost_per_sample: {transmission_cost_per_sample}")


    server = Server.Server(num_users=num_users, num_classes=num_classes, lr=lr, model_name=model_name, wm_config=wm_config, log_dir=logUtil.getDir(),grad_clip=grad_clip, model=top_model)

    if wm_config is not None:
        server.save_wm_kwargs(logUtil.getDir())

    dict_users = DatasetUtil.dataset_iid(train_dataset, num_users)
    dict_users_test = DatasetUtil.dataset_iid(test_dataset, num_users)
    if noniid_mode:
        dict_users = DatasetUtil.dataset_noniid(num_classes=num_classes, dataset=train_dataset,num_users=num_users,beta=noniid_beta)



    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    net_globe_client.to(device)

    # if torch.cuda.device_count() > 1:
    #     net_globe_client = nn.DataParallel(net_globe_client)

    net_globe_client.train()

    w_globe_client = net_globe_client.state_dict()

    clients = []
    for idx in range(num_users):
        clients.append(Client.Client(net_client_model=net_globe_client, server=server, idx=idx, lr=lr, device=device,
                                     local_epoch=local_epochs, dp_epsilon=dp_epsilon, batch_size=batch_size,
                                     dataset_train=train_dataset, dataset_test=test_dataset, indices_train=dict_users[idx],
                                     indices_test=dict_users_test[idx], backdoor_flag=0, grad_clip=grad_clip, rho=rho, bd_overlap=bd_overlap))
    backdoor_activation = False

    flexLR = FlexLearningRate(total_epochs=epochs, initial_lr=lr, decline_multiplier=flex_lr_gamma, num_lr_stages=flex_lr_stages,
                              cos_annealing_multiplier=cos_annealing_multiplier)


    # codes for pre-train to reinforce client-side backdoor
    # pre_train_epochs = 10
    # idxs = [i for i in range(num_users)]
    # for idx in idxs:
    #     clients[idx].activate_backdoor(backdoor_type=backdoor_type)
    # backdoor_activation = True
    #
    # for epoch in range(pre_train_epochs):
    #     w_locals_client = []
    #     current_lr = lr
    #     for idx in idxs:
    #         clients[idx].adjust_learning_rate(current_lr)
    #         w_client = clients[idx].pre_train_backdoor(net=copy.deepcopy(net_globe_client))
    #         w_locals_client.append(copy.deepcopy(w_client))
    #
    #     w_globe_client = Server.fedAvg(w_locals_client)
    #     net_globe_client.load_state_dict(w_globe_client)
    #
    #     for idx in idxs:
    #         clients[idx].evaluate_backdoor_client(net=copy.deepcopy(net_globe_client), ell=epoch)

    # for idx in idxs:
    #     clients[idx].local_ep = local_epochs


    # Manually reset indicators before training process
    # server.l_epoch_check = False
    # server.l_epoch_check_backdoor = False
    # server.fed_check = False
    # server.fed_check_backdoor = False



    for iter in range(epochs):
        m = max(int(frac * num_users), 1)
        idxs_users = np.random.choice(range(num_users), m, replace=False)
        w_locals_client = []

        print('',flush=True)
        current_lr = flexLR.get_lr(iter)
        print("-----------------------------------------------------------")
        print("------ FedServer: Federation process at Client-Side ------- ")
        print(f"-------------------- Epoch: {iter} ----------------------- ")
        print("-----------------------------------------------------------")

        # Clean train before reaching the wm_engage epoch
        if iter < wm_engage * epochs:
        # elif iter < wm_engage * epochs:
            for idx in tqdm(idxs_users, desc=f"Epoch: {iter}, lr: {current_lr}"):
                clients[idx].adjust_learning_rate(current_lr)

                w_client = clients[idx].train(net=copy.deepcopy(net_globe_client))
                w_locals_client.append(copy.deepcopy(w_client))

                clients[idx].evaluate(net=copy.deepcopy(net_globe_client), ell=iter)

        # Server-side watermark engaging
        elif iter < backdoor_engage * epochs:
            for idx in tqdm(idxs_users, desc=f"Epoch: {iter}, lr: {current_lr}"):
                clients[idx].adjust_learning_rate(current_lr)

                w_client = clients[idx].train_with_wm(net=copy.deepcopy(net_globe_client))
                w_locals_client.append(copy.deepcopy(w_client))

                clients[idx].evaluate(net=copy.deepcopy(net_globe_client), ell=iter)

        # Client-side watermark engaging
        else:
            # Initialize clients' backdoors if they were not initialized before
            if not backdoor_activation:
                for idx in idxs_users:
                    clients[idx].activate_backdoor(backdoor_type=backdoor_type)
                    # clients[idx].save_triggered_dataset(log_dir=logUtil.getDir())
                backdoor_activation = True

            for idx in tqdm(idxs_users, desc=f"Epoch: {iter}, lr: {current_lr}"):
                clients[idx].adjust_learning_rate(current_lr)
                w_client = clients[idx].train_with_wm(net=copy.deepcopy(net_globe_client))
                w_locals_client.append(copy.deepcopy(w_client))

                clients[idx].evaluate(net=copy.deepcopy(net_globe_client), ell=iter)
                # clients[idx].evaluate_backdoor_client(net=copy.deepcopy(net_globe_client),ell=iter)

        # Call FedAvg and update the clients' global bottom model when one global epoch finished
        w_globe_client = Server.fedAvg(w_locals_client)
        net_globe_client.load_state_dict(w_globe_client)

        loss, wmr = server.get_wmr()
        print(f"[FedServer] Feature WM: Loss - {loss:.4f}, Wmr - {wmr:.4f}")

        # Evaluate clients backdoors after FedAvg
        if backdoor_activation:
            for idx in idxs_users:
                # clients[idx].evaluate_backdoor_client(net=copy.deepcopy(net_globe_client), ell=iter)
                clients[idx].evaluate_backdoor_client(net=net_globe_client, ell=iter)

    print("Training and Evaluation completed!\n")

    server_time_cost = 0.0
    client_time_cost = 0.0
    for i in range(num_users):
        server_time_cost += clients[i].get_server_time_cost()
        client_time_cost += clients[i].get_client_time_cost()

    server_time_cost /= num_users
    client_time_cost /= num_users

    sys.stdout = sys.__stdout__

    end_time = time.time()

    if backdoor_type == 0:
        backdoor_type_str = "No Backdoor Embedded"
    elif backdoor_type == 1:
        backdoor_type_str = "Random Noise"
    else:
        backdoor_type_str = "Improved C-Pattern"

    _, wm_theta_f = server.get_wm_loss_and_wmr()
    wm_theta_b = server.get_theta_b()
    main_task_acc = server.get_main_task_acc()

    hyperparameters_dict = {
        'model_name': model_name,
        'split_point': split_point,
        'dataset': dataset_name,
        'non-iid': noniid_mode,
        'num_users': num_users,
        'global_epochs': epochs,
        'local_epochs': local_epochs,
        'init_lr': lr,
        'lr_decay_stages': flex_lr_stages,
        'frac': frac,
        'wm_engage': wm_engage,
        'backdoor_engage': backdoor_engage,
        'main_task_acc': main_task_acc,
        'wm_theta_f': wm_theta_f,
        'wm_theta_b': wm_theta_b,
        'bd_intensity':bd_intensity,
        'bd_wm_rho': rho,
        'bd_overlap': bd_overlap,
        'console_log_file':  f"{logUtil.getDir()}{dataset_name}_{now_str}_{epochs}_{wm_engage}_{backdoor_engage}.log",
        'running_time': end_time - start_time,
        'avg_epoch_time': (end_time - start_time) / float(epochs),
        'backdoor_type': backdoor_type_str,
        'server_time_cost_per_epoch': server_time_cost / float(epochs),
        'client_time_cost_per_epoch': client_time_cost / float(epochs),
        'transmission_cost_per_sample': transmission_cost_per_sample,
        'transmission_cost_per_batch': transmission_cost_per_sample * batch_size,
    }

    if dp_epsilon != 0:
        hyperparameters_dict['dp_epsilon'] = args.epsilon
    hyperparameters_dict['global_bottom_model'] = logUtil.save_trained_model(net_globe_client, "global_bottom_model")
    hyperparameters_dict['global_top_model'] = logUtil.save_trained_model(server.get_server_global_top_model(), "global_top_model")
    if noniid_mode:
        hyperparameters_dict['non-iid_beta'] = noniid_beta


    main_task_result_file_name = server.save_main_task_result()
    main_task_detail_result_file_name= server.save_main_task_detail_result()
    backdoor_detail_result_file_name = server.save_backdoor_detail_result()
    backdoor_result_file_name = server.save_backdoor_result()

    hyperparameters_dict['main_task_brief_result'] = main_task_result_file_name
    hyperparameters_dict['main_task_detail_result'] = main_task_detail_result_file_name
    hyperparameters_dict['backdoor_detail_result'] = backdoor_detail_result_file_name
    hyperparameters_dict['backdoor_brief_result'] = backdoor_result_file_name

    logUtil.saveHyperParameters(hyperparameters_dict)

if __name__ == '__main__':

    dirCheck('./logs')
    args = parser_args()
    poisoned_train(args)