FairBADS/fairdata.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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
import os
import pickle
from typing import Dict, Optional

import numpy as np
import torch
from pytorch_lightning import LightningDataModule
from torch.utils.data import DataLoader, Subset, TensorDataset

from utils import add_label_bias, count_label_flips


class FairDataModule(LightningDataModule):
    def __init__(self,
                 data_location,
                 dataset_name,
                 meta_ratio = 0.1,
                 bias_amount=0.4,
                 batch_size=512,
                 num_workers = 0):
        super().__init__()
        self.data_location = data_location
        self.dataset_name = dataset_name
        self.batch_size = batch_size
        self.meta_ratio = meta_ratio
        self.num_workers = num_workers
        self.bias_amount = bias_amount
        self.theta_dict = {'theta_0_p':bias_amount,
                           'theta_0_m':0,
                           'theta_1_p':0,
                           'theta_1_m':bias_amount}

        self.train_file_name = os.path.join(data_location, dataset_name, 'train.pkl')
        self.test_file_name = os.path.join(data_location, dataset_name, 'test.pkl')

    def setup(self, stage= None):
        with open(self.train_file_name, "rb") as f:
            train_raw_data = pickle.load(f)
        x_train, s_train, y_train = train_raw_data['x'], train_raw_data['s'], train_raw_data['y']

        with open(self.test_file_name, "rb") as f:
            test_raw_data = pickle.load(f)
        x_test, s_test, y_test = test_raw_data['x'], test_raw_data['s'], test_raw_data['y']

        # convert to tensor
        x_train = torch.tensor(x_train, dtype=torch.float32)
        s_train = torch.tensor(s_train.squeeze(-1), dtype=torch.long)
        y_train = torch.tensor(y_train.squeeze(-1), dtype=torch.long)

        # divide meta set
        total_indices = np.arange(len(x_train))
        test_indicies = np.arange(len(x_test))
        np.random.shuffle(total_indices)

        n_meta = max(1, int(self.meta_ratio * len(x_train)))
        meta_indices = total_indices[:n_meta]
        train_indices = total_indices[n_meta:]

        self.meta_dataset = TensorDataset(
            x_train[meta_indices],
            y_train[meta_indices],
            s_train[meta_indices],
            torch.tensor(meta_indices, dtype=torch.long)
        )

        if self.bias_amount != 0:
            print("BIAS ADDED")
            y_train_biased = add_label_bias(y_train[train_indices], s_train[train_indices], self.theta_dict)
        else:
            print("ASSUME NO BIAS")
            y_train_biased = y_train[train_indices]

        self.train_dataset = TensorDataset(
            x_train[train_indices],
            y_train_biased,
            s_train[train_indices],
            torch.tensor(train_indices, dtype=torch.long)
        )

        count_label_flips(y_train[train_indices],y_train_biased)

        self.test_dataset = TensorDataset(
            torch.tensor(x_test, dtype=torch.float32),
            torch.tensor(y_test.squeeze(-1), dtype=torch.long),
            torch.tensor(s_test.squeeze(-1), dtype=torch.long),
            torch.tensor(test_indicies, dtype=torch.long)

        )

        s_train_subset = s_train[train_indices]

        unique_groups, group_counts = torch.unique(s_train_subset, return_counts=True)
        total_train = len(s_train_subset)
        group_distribution = {int(g.item()): c.item() / total_train for g, c in zip(unique_groups, group_counts)}

        self.global_group_distribution = group_distribution

        self.image_size = x_train[0].shape


    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size)
    def val_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size)


class DualLoader:
    def __init__(self, train_loader, meta_dataset):
        self.train_loader = train_loader
        self.meta_dataset = meta_dataset

        x_m_list, y_m_list, idx_m_list, s_m_list = [], [], [], []
        for i in range(len(self.meta_dataset)):
            x, y, s, idx = self.meta_dataset[i]
            x_m_list.append(x)
            s_m_list.append(s)
            y_m_list.append(y)
            idx_m_list.append(idx)

        self.x_m = torch.stack(x_m_list)
        self.s_m = torch.tensor(s_m_list)
        self.y_m = torch.tensor(y_m_list)
        self.idx_m = torch.tensor(idx_m_list)

    def __iter__(self):
        self.train_iter = iter(self.train_loader)
        return self

    def __next__(self):
        try:
            train_batch = next(self.train_iter)
        except StopIteration:
            raise StopIteration

        x_t, y_t, s_t, idx_t = train_batch

        return {
            'train': (x_t, y_t, s_t, idx_t),
            'meta': (self.x_m, self.y_m, self.s_m, self.idx_m)
        }

    def __len__(self):
        return len(self.train_loader)