federated_finetuning.py

import argparse
import copy
import os
import random
import timm
import torch
import numpy as np
import util.lr_decay as lrd

from dataset_classes.csi_sensing import CSISensingDataset
from dataset_classes.radio_sig import RadioSignal
from pathlib import Path
from federated_finetuning.data_partition import create_lda_partitions
from federated_finetuning.fl_tools import set_seed, train_one_epoch, aggregate_weights, evaluate
from compression_step.helper_funcs import forward
from collections import Counter
from sklearn.model_selection import StratifiedShuffleSplit
from torch.utils.data import DataLoader, Subset

def main(args):


    NUM_CLIENTS = args.num_clients
    ROUNDS = args.rounds
    LOCAL_EPOCHS = args.local_epochs
    BATCH_SIZE = args.batch_size
    NUM_WORKERS = args.num_workers
    LR = args.lr
    WEIGHT_DECAY = args.weight_decay
    LAYER_DECAY = args.layer_decay
    SEED = args.seed
    log_dir = args.log_dir
    CARB_TRACK = False

    set_seed(SEED)

    if args.task == 'sensing':
       # Load dataset
        dataset_train = CSISensingDataset(Path('/home/ict317-3/Mohammad/mae/fine-tuning_datasets/NTU-Fi_HAR/train'), downsampled=False)
        dataset_val = CSISensingDataset(Path('/home/ict317-3/Mohammad/mae/fine-tuning_datasets/NTU-Fi_HAR/test'), downsampled=False)
    
    elif args.task == 'radio':

        dataset = RadioSignal(Path('fine-tuning_datasets/radio_sig_identification/train')) 
        all_labels = np.array([dataset[i][1] for i in range(len(dataset))])  # Get labels

        # Perform stratified train/val split
        splitter = StratifiedShuffleSplit(n_splits=1, train_size=0.9, test_size=0.1, random_state=SEED)
        train_idx, val_idx = next(splitter.split(range(len(dataset)), all_labels))

        dataset_train = Subset(dataset, train_idx)
        dataset_val = Subset(dataset, val_idx)
        
    # elif args.task == '5g':
        
    x = torch.from_numpy(np.array([sample[0] for sample in dataset_train]))
    y = torch.from_numpy(np.array([sample[1] for sample in dataset_train]))

    # Pack into a tuple
    xy = (x, y)
    if args.partitioning == 'iid':
        alpha = 1e4

    # Sensing dataset is well balanced so we can set alpha to a very small value to simulate non-iid
    elif args.task == 'sensing':
        alpha = 0.1

    # Radio signal dataset is not balanced well so we can have to set alpha value less severly for non-iid
    elif args.task == 'radio':
        alpha = 0.5

    datasets = create_lda_partitions(
        dataset=xy,
        num_partitions=NUM_CLIENTS,
        concentration=alpha,
        accept_imbalanced=True,
        seed=SEED,
    )
        
    # Create client DataLoaders
    client_loaders = []
    for client_id in range(NUM_CLIENTS):

        x_arr, y_arr = datasets[0][client_id]

        subset = [(torch.tensor(x_i, dtype=torch.float32),
                torch.tensor(y_i, dtype=torch.long)) for x_i, y_i in zip(x_arr, y_arr)]

        loader = DataLoader(
            subset, shuffle=True,
            batch_size=BATCH_SIZE,
            num_workers=NUM_WORKERS,
            pin_memory=True,
            drop_last=True
        )
        client_loaders.append(loader)

    # Validation DataLoader (Remains Unchanged)
    data_loader_val = DataLoader(
        dataset_val,
        batch_size=BATCH_SIZE,
        num_workers=NUM_WORKERS,
        pin_memory=True,
        drop_last=False
    )

    # Verify label distribution in each client dataset
    for i, loader in enumerate(client_loaders):
        client_labels = []
        for batch in loader:
            _, labels = batch  # Assuming dataset returns (data, label)
            client_labels.extend(labels.numpy())  # Convert to list
        
        client_label_counts = Counter(client_labels)
        
        print(f"\nClient {i+1} Dataset Label Distribution:")
        for label, count in sorted(client_label_counts.items()):
            print(f"Label {label}: {count} samples")

    seen = set()
    repeated = False

    for i in range(NUM_CLIENTS):
        for sample in datasets[0][i][0]:
            # Convert the array to bytes (hashable)
            sample_bytes = sample.tobytes()
            if sample_bytes in seen:
                repeated = True
                break
            seen.add(sample_bytes)

    if repeated:
        print("\nThere is repeated samples!")
    else:
        print("\nAll samples are unique.")
    
    # Get the sizes of each client's dataset
    client_sizes = []
    for i in range(len(client_loaders)):
        client_sizes.append(len(client_loaders[i].dataset))

    # Global model initialization
    global_model = torch.load(args.pruned_model_path, weights_only=False)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    global_model.to(device)

    for m in global_model.modules():
            if isinstance(m, timm.models.vision_transformer.Attention):
                m.forward = forward.__get__(m, timm.models.vision_transformer.Attention)

    # Freeze all transformer blocks
    for param in global_model.blocks.parameters():
        param.requires_grad = False

    criterion = torch.nn.CrossEntropyLoss()

    acc = 0
    NUM_SELECTED_CLIENTS = 5  # Number of clients that share weights per round

    os.makedirs(log_dir, exist_ok=True)
    model_path = os.path.join(log_dir, "fed_avg_model.pth")

    for round in range(ROUNDS):  
        print(f"Round {round+1}: Training clients...")

        # Randomly select clients for this round
        selected_clients = random.sample(range(NUM_CLIENTS), NUM_SELECTED_CLIENTS)
        print(f"Selected Clients: {selected_clients}")

        client_weights = []
        client_losses = []
        
        for i in selected_clients:  # Train only selected clients
            local_model = copy.deepcopy(global_model)  # Each client starts from the latest global model
            param_groups = lrd.param_groups_lrd(local_model, WEIGHT_DECAY, layer_decay=LAYER_DECAY)  
            optimizer = torch.optim.AdamW(param_groups, lr=LR)
            
            # Train client and get updated weights & loss
            client_state, client_loss = train_one_epoch(local_model, criterion, client_loaders[i], optimizer, device, LOCAL_EPOCHS, CARB_TRACK, log_dir, f"Client_{i}")
            client_weights.append(client_state)
            client_losses.append(client_loss)

        # Aggregate only selected clients' weights and update the global model
        if client_weights:
            round_client_sizes = [client_sizes[j] for j in selected_clients]
            new_global_weights = aggregate_weights(client_weights, round_client_sizes) #aggregate_weights(client_weights)
            global_model.load_state_dict(new_global_weights)

        # Compute average client loss
        avg_client_loss = sum(client_losses) / len(client_losses)
        print(f"Round {round+1}: Average Client Loss = {avg_client_loss:.4f}")

        # Evaluate global model on validation set
        test_stats = evaluate(data_loader_val, global_model, criterion, device)
        print(f"Round {round+1}: Global Model Accuracy = {test_stats['acc1']:.2f}%")

        if acc < test_stats['acc1']:
            acc = test_stats['acc1']
            torch.save(global_model, model_path)
            print('A new better model has been saved ... \n')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Federated Learning with Pruned ViT")

    # Federated learning core config
    parser.add_argument('--num_clients', type=int, default=10,
                        help='Total number of clients in the federated learning setup.')
    parser.add_argument('--num_selected_clients', type=int, default=5,
                        help='Number of clients selected per communication round.')
    parser.add_argument('--rounds', type=int, default=50,
                        help='Number of communication rounds.')
    parser.add_argument('--local_epochs', type=int, default=10,
                        help='Number of local epochs per client per round.')

    # Data loading
    parser.add_argument('--batch_size', type=int, default=64,
                        help='Local training batch size.')
    parser.add_argument('--num_workers', type=int, default=8,
                        help='Number of workers for DataLoader.')

    # Optimization
    parser.add_argument('--lr', type=float, default=1e-3,
                        help='Learning rate for local client training.')
    parser.add_argument('--weight_decay', type=float, default=0.05,
                        help='Weight decay for the optimizer.')
    parser.add_argument('--layer_decay', type=float, default=0.75,
                        help='Layer-wise learning rate decay.')

    # Seed
    parser.add_argument('--seed', type=int, default=22,
                        help='Random seed for reproducibility.')

    # Dataset partitioning
    parser.add_argument('--partitioning', type=str, choices=['iid', 'non-iid'], default='iid',
                        help='Data partitioning strategy across clients: "iid" or "non-iid" (LDA).')
    
    # Paths
    parser.add_argument('--pruned_model_path', type=str, required=True, default='/home/ict317-3/Mohammad/mae/CFFM/pruning_models_small_ViT/pruned_ViT_has/pruned_ViT_with_ratio_50.00_two_layer_head%',
                        help='Path to the initial pruned ViT model.')
    parser.add_argument('--log_dir', type=str, default='/home/ict317-3/Mohammad/mae/has_output_dir',
                        help='Directory to save logs and the best model.')
    
    parser.add_argument('--task', type=str, default='sensing',
                        help='The task that the model will perform. The tasks are: Human Activity Sensing, Radio Signal Identification, 5G Positioning')
    
    args = parser.parse_args()
    main(args)