utils.py

from collections.abc import Mapping
from functools import reduce
import math
import os
import random
import sys
import numpy as np
import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F

from client_opt.fedsophia import SophiaG
from client_opt.ltda import LTDA

def set_random_seeds(seed_value=0, full_determinism=False):
    """Set seed for reproducibility."""
    torch.manual_seed(seed_value)
    torch.cuda.manual_seed_all(seed_value)
    random.seed(seed_value)
    np.random.seed(seed_value)
    if full_determinism:
        os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
        torch.use_deterministic_algorithms(True)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

def get_optimizer(parameters, config):
    """Create an optimizer for training."""
    if config.client_opt == 'sgd':
        optimizer = optim.SGD(parameters, lr=config.lr, momentum=config.momentum,
                              dampening=config.dampening, weight_decay=config.weight_decay, nesterov=config.nesterov)
    elif config.client_opt == 'adam':
        optimizer = optim.Adam(parameters, lr=config.lr, betas=(config.adam_beta1, config.adam_beta2), eps=config.adam_eps,
                               weight_decay=config.weight_decay, amsgrad=config.adam_amsgrad)
    elif config.client_opt == 'adamw':
        optimizer = optim.AdamW(parameters, lr=config.lr, betas=(config.adam_beta1, config.adam_beta2), eps=config.adam_eps,
                                weight_decay=config.weight_decay, amsgrad=config.adam_amsgrad)
    elif config.client_opt == 'sophia':
        optimizer = SophiaG(parameters, lr=config.lr, betas=(config.fedsophia_beta1, config.fedsophia_beta2), rho=config.fedsophia_rho,
                            weight_decay=config.weight_decay, eps=config.fedsophia_eps)
    elif config.client_opt == 'ltda':
        optimizer = LTDA(parameters, lr=config.lr, rho=config.ltda_rho, eps=config.ltda_eps,
                        weight_decay=config.weight_decay, max_grad_norm=config.max_grad_norm)
    else:
        raise ValueError(f"Unsupported optimizer: {config.client_opt}")
    return optimizer

def get_scheduler(optimizer, step_per_epoch, config):
    """Create a learning rate scheduler."""
    if config.lr_scheduler == 'StepLR':
        scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=config.step_size, gamma=config.gamma)
    elif config.lr_scheduler == 'ExponentialLR':
        scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=config.gamma)
    elif config.lr_scheduler == "linear":
        if config.inner_loop > 0:
            num_training_steps = config.communication_rounds * config.inner_loop
        else:
            num_training_steps = config.communication_rounds * step_per_epoch * config.inner_epoch
        num_warmup_steps = math.ceil(num_training_steps * config.warmup_ratio) if config.warmup_steps < 0 else config.warmup_steps
        def lr_lambda(current_step: int):
            if current_step < num_warmup_steps:
                return float(current_step) / float(max(1, num_warmup_steps))
            return max(
                0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
            )
        scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)
    else:
        raise ValueError(f"Unsupported scheduler: {config.scheduler}")
    return scheduler

def get_data_loader(dataset, batch_size, seed, collate_fn, shuffle=True, drop_last=True, sequential=False):
    """Create a data loader for training or evaluation."""
    if sequential:
        g = torch.Generator()
        g.manual_seed(seed)
    else:
        g = None
    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn,
                                       pin_memory=True, drop_last=drop_last, worker_init_fn=seed_worker, generator=g)

def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

def reduce_non_diag(cov_mat, a):
    diag_weight = torch.diag(torch.ones(cov_mat.size(0)) - a).to(cov_mat.device)
    non_diag_weight = torch.zeros_like(diag_weight).fill_(a)
    weight = diag_weight + non_diag_weight
    ret = cov_mat * weight
    return ret

def atleast_1d(tensor_or_array):
    if isinstance(tensor_or_array, torch.Tensor):
        if hasattr(torch, "atleast_1d"):
            tensor_or_array = torch.atleast_1d(tensor_or_array)
        elif tensor_or_array.ndim < 1:
            tensor_or_array = tensor_or_array[None]
    else:
        tensor_or_array = np.atleast_1d(tensor_or_array)
    return tensor_or_array


def torch_pad_and_concatenate(tensor1, tensor2, padding_index=-100):
    """Concatenates `tensor1` and `tensor2` on first axis, applying padding on the second if necessary."""
    tensor1 = atleast_1d(tensor1)
    tensor2 = atleast_1d(tensor2)

    if len(tensor1.shape) == 1 or tensor1.shape[1] == tensor2.shape[1]:
        return torch.cat((tensor1, tensor2), dim=0)

    # Let's figure out the new shape
    new_shape = (tensor1.shape[0] + tensor2.shape[0], max(tensor1.shape[1], tensor2.shape[1])) + tensor1.shape[2:]

    # Now let's fill the result tensor
    result = tensor1.new_full(new_shape, padding_index)
    result[: tensor1.shape[0], : tensor1.shape[1]] = tensor1
    result[tensor1.shape[0] :, : tensor2.shape[1]] = tensor2
    return result

def distribute_data_dirichlet(
    targets, non_iid_alpha, n_workers, num_samples=None, seed=0
):
    """Code adapted from Tao Lin (partition_data.py)"""
    random_state = np.random.RandomState(seed=seed)
    num_classes = len(np.unique(targets))
    indices2targets = np.array(list(enumerate(targets)))
    idx_batch = []
    # rebuild _targets.
    _targets = np.array(indices2targets)
    _targets_size = len(_targets)
    # use auxi_workers for this subset targets.
    _n_workers = n_workers
    # get the corresponding idx_batch.
    min_size = 0
    sizes = [0 for _ in range(_n_workers)]
    while True:
        _idx_batch = [[] for _ in range(_n_workers)]
        for _class in range(num_classes):
            # get the corresponding indices in the original 'targets' list.
            idx_class = np.where(_targets[:, 1] == _class)[0]
            idx_class = _targets[idx_class, 0]
            # sampling.
            try:
                proportions = random_state.dirichlet(
                    np.repeat(non_iid_alpha, _n_workers)
                )
                if num_samples is not None:
                    proportions = proportions * np.array(num_samples) / sum(num_samples)
                # balance
                if num_samples is None:
                    proportions = np.array(
                        [
                            p * (len(idx_j) < _targets_size / _n_workers)
                            for p, idx_j in zip(proportions, _idx_batch)
                        ]
                    )
                else:
                    proportions = np.array(
                        [
                            p * (len(idx_j) < num_sample)
                            for p, idx_j, num_sample in zip(proportions, _idx_batch, num_samples)
                        ]
                    )
                proportions = proportions / proportions.sum()
                proportions = (np.cumsum(proportions) * len(idx_class)).astype(int)[
                    :-1
                ]
                _idx_batch = [
                    idx_j + idx.tolist()
                    for idx_j, idx in zip(
                        _idx_batch, np.split(idx_class, proportions)
                    )
                ]
                sizes = [len(idx_j) for idx_j in _idx_batch]
                min_size = min([_size for _size in sizes])
            except ZeroDivisionError:
                pass
        if ((num_samples is None and min_size >= int(0.5 * _targets_size / _n_workers)) or
            (num_samples is not None and all([_size >= 0.5 * num_sample for _size, num_sample in zip(sizes, num_samples)]))):
            break
    idx_batch += _idx_batch
        
    return idx_batch

def im2col_2d(x: torch.Tensor, conv2d: nn.Module):
    if x.ndim != 4:  # n x c x h_in x w_in
        raise ValueError(f'x.ndim has to be 4. Got {x.ndim}.')
    if not isinstance(conv2d, (nn.Conv2d, nn.ConvTranspose2d)):
        raise TypeError(f'conv2d has to be {nn.Conv2d} or {nn.ConvTranspose2d}. Got {type(conv2d)}.')
    if conv2d.dilation != (1, 1):
        return im2col_2d_slow(x, conv2d)

    ph, pw = conv2d.padding if conv2d.padding != 'valid' else (0, 0)
    kh, kw = conv2d.kernel_size
    sy, sx = conv2d.stride
    if ph + pw > 0:
        x = F.pad(x, (pw, pw, ph, ph)).data
    x = x.unfold(2, kh, sy)  # n x c x h_out x w_in x kh
    x = x.unfold(3, kw, sx)  # n x c x h_out x w_out x kh x kw
    x = x.permute(0, 1, 4, 5, 2,
                  3).contiguous()  # n x c x kh x kw x h_out x w_out
    x = x.view(x.size(0),
               x.size(1) * x.size(2) * x.size(3),
               x.size(4) * x.size(5))  # n x c(kh)(kw) x (h_out)(w_out)
    return x

def im2col_2d_slow(x: torch.Tensor, conv2d: nn.Module):
    if x.ndim != 4:  # n x c x h_in x w_in
        raise ValueError(f'x.ndim has to be 4. Got {x.ndim}.')
    if not isinstance(conv2d, (nn.Conv2d, nn.ConvTranspose2d)):
        raise TypeError(f'conv2d has to be {nn.Conv2d} or {nn.ConvTranspose2d}. Got {type(conv2d)}.')

    padding = conv2d.padding if conv2d.padding != 'valid' else (0, 0)
    # n x c(k_h)(k_w) x (h_out)(w_out)
    Mx = F.unfold(x,
                  conv2d.kernel_size,
                  dilation=conv2d.dilation,
                  padding=padding,
                  stride=conv2d.stride)

    return Mx

def get_module_by_name(module, access_string):
    """Retrieve a module nested in another by its access string.

    Works even when there is a Sequential in the module.
    """
    names = access_string.split(sep='.')
    if '' in names:
        return module
    return reduce(getattr, names, module)

def prepare_input(data, device):
    if isinstance(data, Mapping):
        return type(data)({k: prepare_input(v, device) for k, v in data.items()})
    elif isinstance(data, (tuple, list)):
        return type(data)(prepare_input(v, device) for v in data)
    elif isinstance(data, torch.Tensor):
        return data.to(device)
    return data

def prepare_inputs(inputs, device, model_name):
    inputs = prepare_input(inputs, device)
    if len(inputs) == 0:
        raise ValueError("The batch received was empty, your model won't be able to train on it.")
    if model_name == "roberta-base":
        new_inputs = {}
        for k in ["input_ids", "attention_mask", "labels", "dataset"]:
            if k in inputs:
                new_inputs[k] = inputs[k]
        labels = None
        if "labels" in new_inputs:
            labels = new_inputs["labels"]
        return new_inputs, labels
    else:
        return {'x': inputs[0], 'labels': inputs[1]}, inputs[1]

def get_size(obj, seen=None):
    """Recursively finds size of objects"""
    size = sys.getsizeof(obj)
    if seen is None:
        seen = set()
    obj_id = id(obj)
    if obj_id in seen:
        return 0
    # Important mark as seen *before* entering recursion to gracefully handle
    # self-referential objects
    seen.add(obj_id)
    if isinstance(obj, dict):
        size += sum([get_size(v, seen) for v in obj.values()])
        size += sum([get_size(k, seen) for k in obj.keys()])
    elif hasattr(obj, '__dict__'):
        size += get_size(obj.__dict__, seen)
    elif hasattr(obj, '__iter__') and not isinstance(obj, (str, bytes, bytearray)):
        size += sum([get_size(i, seen) for i in obj])
    return size

def select_participants(num_clients, num_participants):
    # Randomly select 'num_participants' nodes to participate in the training round
    return sorted(torch.randperm(num_clients)[:num_participants].tolist())