samoo/dataset.py at master · proteekroy/samoo · 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
from torch.utils.data import Dataset


class DatasetPrivacy(Dataset):
    def __init__(self, index_tensor, data_tensor, target_tensor, sensitive_tensor):
        Dataset.__init__(self)
        assert data_tensor.size(0) == target_tensor.size(0)
        self.data_tensor = data_tensor
        self.target_tensor = target_tensor
        self.sensitive_tensor = sensitive_tensor
        self.index_tensor = index_tensor

    def __getitem__(self, index):
        return self.index_tensor[index], self.data_tensor[index], self.target_tensor[index], self.sensitive_tensor[
            index]

    def __len__(self):
        return self.data_tensor.size(0)


class DatasetOpt(Dataset):
    def __init__(self, index_tensor, x, f, g, g_label=None):
        Dataset.__init__(self)
        assert index_tensor == x.size(0)
        assert x.size(0) == f.size(0)
        assert x.size(0) == g.size(0)
        assert x.size(0) == g_label.size(0)

        self.index_tensor = index_tensor
        self.x = x
        self.f = f
        self.g = g
        self.g_label = g_label

    def __getitem__(self, index):
        return self.index_tensor[index], self.x[index], self.f[index], self.g[index], self.g_label[index]

    def __len__(self):
        return self.x.size(0)


class DatasetFunction(Dataset):
    def __init__(self, index_tensor, x, f):
        Dataset.__init__(self)
        assert index_tensor.size(0) == x.size(0)
        assert x.size(0) == f.size(0)

        self.index_tensor = index_tensor
        self.x = x
        self.f = f

    def __getitem__(self, index):
        return self.index_tensor[index], self.x[index], self.f[index]

    def __len__(self):
        return self.x.size(0)


class DatasetConstraint(Dataset):
    def __init__(self, index_tensor, x, g, g_label=None):
        Dataset.__init__(self)
        assert index_tensor.size(0) == x.size(0)
        assert x.size(0) == g.size(0)
        assert x.size(0) == g_label.size(0)

        self.index_tensor = index_tensor
        self.x = x
        self.g = g
        self.g_label = g_label

    def __getitem__(self, index):
        return self.index_tensor[index], self.x[index], self.g[index], self.g_label[index]

    def __len__(self):
        return self.x.size(0)