<filename>test/FedAvgBenchmarks/local_train.py
'''
# File: ptrain.py
# Project: pFedGD
# Created Date: 2022-03-15 11:26
# Author: <NAME> <EMAIL>
# -----
# Last Modified: 2022-04-29 10:32
# Modified By: <NAME> <EMAIL>
# 
# This code is published under the MIT License.
# -----
# HISTORY:
# Date      	By 	Comments
# ----------	---	----------------------------------------------------------
'''


import os
import torch
import torch.nn as nn
import yaml
import sys
import argparse

sys.path.append("../../")
from src.flbase.utils import create_clients_from_existing_ones, setup_clients
from src.flbase.strategies.FedAvg import FedAvgClient
from src.utils import setup_seed, mkdirs, get_datasets, load_from_pkl, save_to_pkl
from torch.utils.data import DataLoader

global wandb_installed
try:
    import wandb
    wandb_installed = True
except ModuleNotFoundError:
    wandb_installed = False

def run(server_pkl_path, use_wandb, **kwargs):
    use_wandb = wandb_installed and use_wandb
    setup_seed(2022)
    saved_server = load_from_pkl(server_pkl_path)
    server_config = saved_server.server_config
    client_config = saved_server.clients_dict[0].client_config
    datasetname =  server_config['dataset']
    config = {'server_config':server_config, 'client_config':client_config}
    num_clients = server_config['num_clients']
   
    run_tag = f"localonly_{server_config['dataset']}_{client_config['model']}_{server_config['partition']}_beta:{server_config['beta']}_num-classes-per-client:{server_config['num_classes_per_client']}_num-clients:{server_config['num_clients']}"
    if use_wandb:
        wandb.init(config=config, name=run_tag)

    criterion = nn.CrossEntropyLoss()

    trainset, testset, _ = get_datasets(server_config['dataset'])
    global_testloader = DataLoader(testset, batch_size=64, shuffle=True)
    # get starting weight (traning from scractch)
    old_clients_dict = setup_clients(FedAvgClient, trainset, None, criterion, num_clients, [client_config for i in range(num_clients)], 'cuda',
                            strategy=server_config['partition'], num_classes=server_config['num_classes'], beta=server_config['beta'], num_classes_per_client=server_config['num_classes_per_client'])
    init_weight = old_clients_dict[0].get_params()                    
    clients_dict = old_clients_dict
    directory = './experiments/'
    mkdirs(directory)
    path = directory + run_tag
    print('results are saved in: ',path)
    for cid in clients_dict.keys():
        print(f"Progress:{cid}/{len(clients_dict)}")
        old_client = saved_server.clients_dict[cid]
        client = clients_dict[cid]
        num_rounds_particiapted = old_client.num_rounds_particiapted
        print(f"Client {cid} perform {num_rounds_particiapted} epochs local training. Now: {len(client.trainset)} {client.label_dist} | Old: {len(old_client.trainset)} {old_client.label_dist} ")
        client.set_params(init_weight)
        for r in range(1, num_rounds_particiapted+1):
            client.training(r, client.client_config['num_epochs'])
            client.testing(r, global_testloader, weights=['labeldist', 'validclass'], personal=True)
            print(f"round{r:3d} |  test acc (labeldist):{client.test_pfl_acc_dict[r]['labeldist'][-1]:.3f} | test acc (validclass):{client.test_pfl_acc_dict[r]['validclass'][-1]:.3f}")
        client.model = None
    save_to_pkl(clients_dict, path + '_final_clients_dict.pkl')


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Test local only.')
    parser.add_argument('--server_pkl_path', default=None, type=str, help='path to saved FedAvg Server')
    parser.add_argument('--use_wandb', default=False, type=lambda x: (str(x).lower() in ['true', '1', 'yes']), help='Use wandb pkg')
    args = parser.parse_args()
    run(args.server_pkl_path, args.use_wandb)
