FairBADS/utils.py at main · co234/FairBADS · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import torch
import numpy as np
from torch.utils import data
import torch.nn as nn
import torch
import pickle
import math
from sklearn.utils import resample
from sklearn import svm
import warnings
import clip
import torch.nn.functional as F
from torchvision import transforms

warnings.filterwarnings("ignore")


def add_label_bias(yclean, rho, theta_dict, seed=1975):
    """
    Add bias to labels based on sensitive attributes.

    theta_0_p: P(Y=+1|Z=-1,A=0)
    theta_0_m: P(Y=-1|Z=+1,A=0)
    theta_1_p: P(Y=+1|Z=-1,A=1)
    theta_1_m: P(Y=-1|Z=+1,A=1)
    """

    t_0_p, t_0_m, t_1_p, t_1_m = theta_dict['theta_0_p'], theta_dict['theta_0_m'], theta_dict['theta_1_p'], theta_dict['theta_1_m']


    g_01 = (rho == 0) & (yclean == 1)
    g_00 = (rho == 0) & (yclean == 0)
    g_11 = (rho == 1) & (yclean == 1)
    g_10 = (rho == 1) & (yclean == 0)

    group = [g_01, g_00, g_11, g_10]
    theta = [t_0_m, t_0_p, t_1_m, t_1_p]
    tilde_y = [0, 1, 0, 1]

    t = yclean.clone()  # Make a copy of yclean

    for i in range(len(group)):
        idxs = group[i].nonzero(as_tuple=True)[0]
        for idx in idxs:
            p = torch.rand(1).item()
            if p < theta[i]:
                t[idx] = tilde_y[i]

    return t


def count_label_flips(yclean, ybiased):
    flips = (yclean != ybiased).sum().item()
    total = len(yclean)
    print(f"Flipped {flips}/{total} labels ({100 * flips / total:.2f}%)")


def compute_rbf_kernel(X, h=None):
    pairwise_dists = torch.cdist(X, X) ** 2
    if h is None:
        h = torch.median(pairwise_dists)
        h = h / torch.log(torch.tensor(X.shape[0], dtype=torch.float32) + 1.0)

    kxy = torch.exp(-pairwise_dists / h)
    dxkxy = -2 / h * (X.unsqueeze(1) - X.unsqueeze(0)) * kxy.unsqueeze(-1)
    dxkxy = dxkxy.sum(dim=1)

    return kxy, dxkxy


def compute_particle_diversity(particles):
    param_vectors = []
    for p in particles:
        θ = torch.cat([param.detach().flatten() for param in p.model.parameters()])
        param_vectors.append(θ)
    param_vectors = torch.stack(param_vectors)
    pdist = torch.cdist(param_vectors, param_vectors)
    diversity = pdist.mean().item()
    return diversity


def compute_clip_meta_loss(x_m, y_m, dataset_name, device):

    model, preprocess = clip.load("RN50", device=device)
    model.eval()

    if dataset_name.lower() == "celeba":
        class_names = ["no heavy makeup", "heavy makeup"]
    else:
        class_names = ["female", "male"]

    text_inputs = torch.cat([clip.tokenize(f"a photo of a {name}") for name in class_names]).to(device)
    resize_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        preprocess.transforms[1]
    ])

    if x_m.shape[1] == 1:
        x_m = x_m.repeat(1, 3, 1, 1)

    x_m = resize_transform(x_m)

    with torch.no_grad():
        image_features = model.encode_image(x_m)
        text_features = model.encode_text(text_inputs)

        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)

        similarity = image_features @ text_features.T  # [B, C]
        clip_probs = similarity.softmax(dim=-1)
        loss = F.cross_entropy(similarity, y_m)

    return clip_probs