-
Notifications
You must be signed in to change notification settings - Fork 1
/
mapping.py
242 lines (214 loc) · 9.22 KB
/
mapping.py
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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
from torch.nn.functional import one_hot
from tqdm import tqdm
import torch.nn.functional as F
import torch
import torch.nn as nn
def label_mapping_base(logits, mapping_sequence):
'''
:param logits: output of the pretrained model
:param mapping_sequence: one-to-one label mapping matrix
:return: predicted results on the target label space
'''
modified_logits = logits[:, mapping_sequence]
return modified_logits
def label_mapping_calculation(logits, mapping_matrix):
'''
:param logits: output of the pretrained model
:param mapping_sequence: real reweight matrix $\omega$
:return: predicted results on the target label space
'''
modified_logits = torch.mm(logits, mapping_matrix)
return modified_logits
def get_freq_distribution(fx, y):
'''
:param fx: the logits output of the pretrained model
:param y: the groud truth labels
:return: frequency distribution matrix of [source predicted labels, target ground truth labels]
'''
fx = one_hot(torch.argmax(fx, dim=-1), num_classes=fx.size(-1))
freq_matrix = [fx[y==i].sum(0).unsqueeze(1) for i in range(len(y.unique()))]
freq_matrix = torch.cat(freq_matrix, dim=1)
return freq_matrix
def greedy_mapping(freq_matrix):
'''
:param freq_matrix: frequency distribution matrix of [source predicted labels, target ground truth labels]
:return: greedy one-to-one mapping results
'''
mapping_matrix = torch.zeros_like(freq_matrix, dtype=int)
freq_matrix_flat = freq_matrix.flatten()
for _ in range(freq_matrix.size(1)):
loc = freq_matrix_flat.argmax().item()
loc = [loc // freq_matrix.size(1), loc % freq_matrix.size(1)]
mapping_matrix[loc[0], loc[1]] = 1
freq_matrix[loc[0]] = -1
if mapping_matrix[:, loc[1]].sum() == 1:
freq_matrix[:, loc[1]] = -1
return mapping_matrix
def one2one_mappnig_matrix(visual_prompt, network, data_loader):
'''
The optimal one-to-one label mapping (for FLM, ILM)
:param visual_prompt: Current input VR
:param network: Pretrained model
:param data_loader: Dataloader for downstream tasks
:return: optimal one-to-one label mapping
'''
device = next(visual_prompt.parameters()).device
if hasattr(network, "eval"):
network.eval()
fx0s = []
ys = []
pbar = tqdm(data_loader, total=len(data_loader), desc=f"FLM", ncols=100) if len(data_loader) > 20 else data_loader
for x, y in pbar:
x, y = x.to(device), y.to(device)
with torch.no_grad():
fx0 = network(visual_prompt(x))
fx0s.append(fx0)
ys.append(y)
fx0s = torch.cat(fx0s).cpu().float()
ys = torch.cat(ys).cpu().int()
freq_matrix = get_freq_distribution(fx0s, ys)
pairs = torch.nonzero(greedy_mapping(freq_matrix))
mapping_sequence = pairs[:, 0][torch.sort(pairs[:, 1]).indices.tolist()]
return mapping_sequence
def blm_reweight_matrix(visual_prompt, network, data_loader, lap=1):
'''
The optimal real reweight mapping matrix (Bayesian-Guided Label Mapping)
:param visual_prompt: Current input VR
:param network: Pretrained model
:param data_loader: Dataloader for downstream tasks
:param lap: laplace smooth factor - $\lambda$
:return: optimal real reweight mapping matrix
'''
device = next(visual_prompt.parameters()).device
if hasattr(network, "eval"):
network.eval()
fx0s = []
ys = []
pbar = tqdm(data_loader, total=len(data_loader), desc=f"BLM", ncols=100) if len(data_loader) > 20 else data_loader
for x, y in pbar:
x, y = x.to(device), y.to(device)
with torch.no_grad():
fx0 = network(visual_prompt(x))
fx0s.append(fx0)
ys.append(y)
fx0s = torch.cat(fx0s).cpu().float()
ys = torch.cat(ys).cpu().int()
freq_matrix = get_freq_distribution(fx0s, ys)
# Begin calculating marginal distribution
classes_sum = torch.sum(freq_matrix, dim=1, keepdim=True)
# Laplace Smoothing
classes_sum = classes_sum + lap * torch.ones_like(classes_sum)
matrix = torch.div(freq_matrix, classes_sum)
# Normalization
target_sum = torch.sum(matrix, dim=0)
norm_matrix = torch.div(matrix, target_sum)
return norm_matrix.to(device)
def blmp_reweight_matrix(visual_prompt, network, data_loader, lap=0, k=3):
'''
The optimal real reweight mapping matrix (Improved Bayesian-Guided Label Mapping)
:param visual_prompt: Current input VR
:param network: Pretrained model
:param data_loader: Dataloader for downstream tasks
:param lap: laplace smooth factor - $\lambda$
:param k: Top k factor
:return: optimal real reweight mapping matrix
'''
device = next(visual_prompt.parameters()).device
if hasattr(network, "eval"):
network.eval()
probs_list = []
ys = []
pbar = tqdm(data_loader, total=len(data_loader), desc=f"BLM+", ncols=100) if len(data_loader) > 20 else data_loader
for x, y in pbar:
x, y = x.to(device), y.to(device)
with torch.no_grad():
fx0 = network(visual_prompt(x))
probabilities = F.softmax(fx0, dim=1)
# Top-k Mechanism
top_values, top_indices = torch.topk(probabilities, k, dim=1)
# Probability Summation Module
result = torch.zeros_like(probabilities)
result.scatter_(1, top_indices, top_values)
probs_list.append(result.cpu().float())
ys.append(y)
probs = torch.cat(probs_list, dim=0)
ys = torch.cat(ys).cpu().int()
matrix = torch.zeros((len(ys.unique()), probs.size(-1)))
indices = ys.view(-1, 1).long().expand(-1, matrix.size(-1))
matrix.scatter_add_(0, indices, probs)
matrix = matrix.t()
# Begin calculating marginal distribution
classes_sum = torch.sum(matrix, dim=1, keepdim=True)
# Laplace Smoothing
classes_sum = classes_sum + lap * torch.ones_like(classes_sum)
matrix = torch.div(matrix, classes_sum)
# Normalization
target_sum = torch.sum(matrix, dim=0)
norm_matrix = torch.div(matrix, target_sum)
return norm_matrix.to(device)
def update_blmp_reweight_matrix(probs, ys, device, lap=0):
'''
Updating the optimal real reweight mapping matrix after the first epoch (Improved Bayesian-Guided Label Mapping)
:param probs: the stored probabilities of the last epoch
:param ys: the stored ground truth labels of the last epoch
:param device: current device 'cpu' or 'cuda:0'
:param lap: laplace smooth factor - $\lambda$
:return: the updated optimal real reweight mapping matrix
'''
matrix = torch.zeros((len(ys.unique()), probs.size(-1)))
indices = ys.view(-1, 1).long().expand(-1, matrix.size(-1))
matrix.scatter_add_(0, indices, probs)
matrix = matrix.t()
# Begin calculating marginal distribution
classes_sum = torch.sum(matrix, dim=1, keepdim=True)
# Laplace Smoothing
classes_sum = classes_sum + lap * torch.ones_like(classes_sum)
matrix = torch.div(matrix, classes_sum)
# Normalization
target_sum = torch.sum(matrix, dim=0)
norm_matrix = torch.div(matrix, target_sum)
return norm_matrix.to(device)
def update_blm_reweight_matrix(fx0s, ys, device, lap=1):
'''
Updating the optimal real reweight mapping matrix after the first epoch (Bayesian-Guided Label Mapping)
:param fx0s: the stored output results of the last epoch
:param ys: the stored ground truth labels of the last epoch
:param device: current device 'cpu' or 'cuda:0'
:param lap: laplace smooth factor - $\lambda$
:return: the updated optimal real reweight mapping matrix
'''
dist_matrix = get_freq_distribution(fx0s, ys)
# Begin calculating marginal distribution
classes_sum = torch.sum(dist_matrix, dim=1, keepdim=True)
# Laplace Smoothing
classes_sum = classes_sum + lap * torch.ones_like(classes_sum)
matrix = torch.div(dist_matrix, classes_sum)
# Normalization
target_sum = torch.sum(matrix, dim=0)
norm_matrix = torch.div(matrix, target_sum)
return norm_matrix.to(device)
def update_one2one_mappnig_matrix(fx0s, ys):
'''
Updating the optimal real reweight mapping matrix after the first epoch (Iterative Label Mapping)
:param fx0s: the stored output results of the last epoch
:param ys: the stored ground truth labels of the last epoch
:return: the updated optimal one-to-one mapping sequence
'''
freq_matrix = get_freq_distribution(fx0s, ys)
pairs = torch.nonzero(greedy_mapping(freq_matrix))
mapping_sequence = pairs[:, 0][torch.sort(pairs[:, 1]).indices.tolist()]
return mapping_sequence
class FTlayer(nn.Module):
def __init__(self, class_num, norm='none'):
super(FTlayer, self).__init__()
self.norm = norm
if self.norm == 'none':
self.linear = nn.Linear(1000, class_num)
else:
self.linear = nn.Linear(1000, class_num, bias=False)
def forward(self, x):
if self.norm == 'sigmoid':
weights = torch.sigmoid(self.linear.weight)
else:
weights = self.linear.weight
return x @ weights.T