utils.py

import os
import json
import csv 
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import torch
from scipy.stats import norm

from PIL import Image

from config import get_config
config, unparsed = get_config()


def denormalize(T, coords):
    return (0.5 * ((coords + 1.0) * T))


class AverageMeter(object):
    """
    Computes and stores the average and
    current value.
    """

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        # correct_k = correct[:k].view(-1).float().sum(0)
        correct_k = correct[:k].reshape(-1, 1).view(-1).float().sum(0)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


def resize_array(x, size):
    # 3D and 4D tensors allowed only
    assert x.ndim in [3, 4], "Only 3D and 4D Tensors allowed!"

    # 4D Tensor
    if x.ndim == 4:
        res = []
        for i in range(x.shape[0]):
            img = array2img(x[i])
            img = img.resize((size, size))
            img = np.asarray(img, dtype='float32')
            img = np.expand_dims(img, axis=0)
            img /= 255.0
            res.append(img)
        res = np.concatenate(res)
        res = np.expand_dims(res, axis=1)
        return res

    # 3D Tensor
    img = array2img(x)
    img = img.resize((size, size))
    res = np.asarray(img, dtype='float32')
    res = np.expand_dims(res, axis=0)
    res /= 255.0
    return res


def img2array(data_path, desired_size=None, expand=False, view=False):
    """
    Util function for loading RGB image into a numpy array.

    Returns array of shape (1, H, W, C).
    """
    img = Image.open(data_path)
    img = img.convert('RGB')
    if desired_size:
        img = img.resize((desired_size[1], desired_size[0]))
    if view:
        img.show()
    x = np.asarray(img, dtype='float32')
    if expand:
        x = np.expand_dims(x, axis=0)
    x /= 255.0
    return x


def array2img(x):
    """
    Util function for converting anumpy array to a PIL img.

    Returns PIL RGB img.
    """
    x = np.asarray(x)
    x = x + max(-np.min(x), 0)
    x_max = np.max(x)
    if x_max != 0:
        x /= x_max
    x *= 255
    return Image.fromarray(x.astype('uint8'), 'RGB')


def prepare_dirs(config):
    for path in [config.ckpt_dir, config.logs_dir]:
        if not os.path.exists(path):
            os.makedirs(path)


def save_config(config):
    model_name = config.save_name
    filename = model_name + '_params.json'
    param_path = os.path.join(config.ckpt_dir, filename)

    print("[*] Model Checkpoint Dir: {}".format(config.ckpt_dir))
    print("[*] Param Path: {}".format(param_path))

    with open(param_path, 'w') as fp:
        json.dump(config.__dict__, fp, indent=4, sort_keys=True)


def gaussian_mechanism(gradient, sigma_g, max_norm, batch_size):
    std_gaussian = (2 * sigma_g + max_norm) / batch_size
    # print(gradient.shape)
    noise = torch.normal(0, std_gaussian, gradient.shape).cuda()
    
    return gradient + noise


def asymmetric_double_exponential_pdf(x, kappa, theta, sigma):
    mul = np.sqrt(2) * kappa / (1 + kappa ** 2)
    new_x = (x - theta) / sigma 
    if new_x >= 0: 
        return mul * np.exp(-np.sqrt(2) * kappa * np.abs(new_x)) 
    else:
        return mul * np.exp((-np.sqrt(2) * np.abs(new_x)) / kappa) 
    
    
def lambda_mapper(p_value, given_epoch=50): 

    tau_l = 0.001 
    tau_u = 0.1
    tau_opt = 0.01
    l = 1.0 
    lambda_max = 100.0

    if given_epoch < 26: 
        return l * lambda_max / 2.0 


    if p_value < tau_l: 
        l = 1.0 
    elif tau_l <= p_value <= tau_opt:
        l = (lambda_max - 1) * (p_value - tau_l) / (tau_opt - tau_l) + 1
    elif tau_opt <= p_value <= tau_u:
        l = -lambda_max * (p_value - tau_opt) / (tau_u - tau_opt) + lambda_max
    else:
        l = 0.0 
    
    return l  


def lambda_mapper_v2(p_value, given_epoch=50): 
    tau_u = 0.1 
    tau_opt = 0.01 

    l = 1.0  
    lambda_max = 25.0 

    if given_epoch < 26: 
        return l * lambda_max # / 2.0 

    if p_value <= tau_opt:
        l = lambda_max 
    elif tau_opt <= p_value <= tau_u:
        l = -lambda_max * (p_value - tau_opt) / (tau_u - tau_opt) + lambda_max
    else:
        l = 0.0 
    
    return l  


def pvalue(klloss, mean, std):
    p_value = norm.cdf(klloss, mean, std) 

    # Compute p-value
    # if klloss < mean:
    #     p_value = cdf_value
    # else:
    #     p_value = 1 - cdf_value 
    
    return p_value # p_value 


def reverse_sigmoid(tau, a, tau_mid): 
    return 1 - 1 / (1 + np.exp(-a * (tau - tau_mid))) 

def sigmoid(x, a, b):
    return 1 / (1 + np.exp(-a * (x - b)))


def smooth_lambda_mapper(p_value, given_epoch=50):
    tau_u = 0.1 
    tau_opt = 0.01 
    tau_mid = (tau_opt + tau_u) / 2 
    a = 100  # 100  # This value can be adjusted for steepness 

    l = 1.0 
    lambda_max = config.total_lambda / (config.model_num - 1)  # 100.0 

    if given_epoch < 26: 
        return 25.0  # l * lambda_max / 4.0 
    
    l = reverse_sigmoid(p_value, a, tau_mid)     
    return l * lambda_max 


def cosine_mapper(cos_value):
    # Define sigmoid parameters
    a = 25  # Adjust this value for the desired sigmoid shape
    b = 0.5  # Adjust this value to shift the sigmoid horizontally 

    # Compute sigmoid values
    l = sigmoid(cos_value, a, b) 
    lambda_max = 100.0 
    return l * lambda_max 


def ce_ratio_mapper(ratio):
    a = 25 
    b = 1.0 
    lambd = sigmoid(ratio, a, b)
    return lambd 


def lambda_mapper_method1(p_value, tau_u, tau_opt, given_epoch=50): 
    l = 1.0 
    lambda_max = config.total_lambda / (config.model_num - 1) 

    if given_epoch < 26: 
        return l * lambda_max # / 2.0 

    if p_value <= tau_opt: 
        l = lambda_max 
    elif tau_opt <= p_value <= tau_u:
        l = -lambda_max * (p_value - tau_opt) / (tau_u - tau_opt) + lambda_max
    else:
        l = 0.0 
    
    return l  


def log_writer(filename, dictionary, type='kl_loss'): 
    if type == 'kl_loss':
        with open(filename, 'a+', newline='') as csvfile:
            csvfile.seek(0)
            fieldnames = ['epoch', 'user1', 'user2', 'batch_idx', 'kl_loss']
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames) 
            
            if not csvfile.read(1):
                writer.writeheader()
            
            writer.writerow(dictionary) 
    
    elif type == 'cos_sim':
        with open(filename, 'a+', newline='') as csvfile:
            csvfile.seek(0)
            fieldnames = ['epoch', 'user1', 'user2', 'batch_idx', 'cos_sim'] 
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames) 
            
            if not csvfile.read(1):
                writer.writeheader()
            
            writer.writerow(dictionary) 
    
    elif type == 'marginal':
        with open(filename, 'a+', newline='') as csvfile:
            csvfile.seek(0)
            fieldnames = ['epoch', 'user1', 'user2', 'batch_idx', 'loss', 'acc'] 
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames) 
            
            if not csvfile.read(1):
                writer.writeheader()
            
            writer.writerow(dictionary) 
    
    elif type == 'lambda':
        with open(filename, 'a+', newline='') as csvfile:
            csvfile.seek(0)
            fieldnames = ['epoch', 'user1', 'user2', 'batch_idx', 'lambda', 'pvalue'] 
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames) 
            
            if not csvfile.read(1):
                writer.writeheader()
            
            writer.writerow(dictionary)