forked from kohpangwei/group_DRO
-
Notifications
You must be signed in to change notification settings - Fork 0
/
loss.py
195 lines (162 loc) · 8.47 KB
/
loss.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
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class LossComputer:
def __init__(self, criterion, is_robust, dataset, alpha=None, gamma=0.1, adj=None, min_var_weight=0, step_size=0.01, normalize_loss=False, btl=False):
self.criterion = criterion
self.is_robust = is_robust
self.gamma = gamma
self.alpha = alpha
self.min_var_weight = min_var_weight
self.step_size = step_size
self.normalize_loss = normalize_loss
self.btl = btl
self.n_groups = dataset.n_groups
self.group_counts = dataset.group_counts().cuda()
self.group_frac = self.group_counts/self.group_counts.sum()
self.group_str = dataset.group_str
if adj is not None:
self.adj = torch.from_numpy(adj).float().cuda()
else:
self.adj = torch.zeros(self.n_groups).float().cuda()
if is_robust:
assert alpha, 'alpha must be specified'
# quantities maintained throughout training
self.adv_probs = torch.ones(self.n_groups).cuda()/self.n_groups
self.exp_avg_loss = torch.zeros(self.n_groups).cuda()
self.exp_avg_initialized = torch.zeros(self.n_groups).byte().cuda()
self.reset_stats()
def loss(self, yhat, y, group_idx=None, is_training=False):
# compute per-sample and per-group losses
per_sample_losses = self.criterion(yhat, y)
group_loss, group_count = self.compute_group_avg(per_sample_losses, group_idx)
group_acc, group_count = self.compute_group_avg((torch.argmax(yhat,1)==y).float(), group_idx)
# update historical losses
self.update_exp_avg_loss(group_loss, group_count)
# compute overall loss
if self.is_robust and not self.btl:
actual_loss, weights = self.compute_robust_loss(group_loss, group_count)
elif self.is_robust and self.btl:
actual_loss, weights = self.compute_robust_loss_btl(group_loss, group_count)
else:
actual_loss = per_sample_losses.mean()
weights = None
# update stats
self.update_stats(actual_loss, group_loss, group_acc, group_count, weights)
return actual_loss
def compute_robust_loss(self, group_loss, group_count):
adjusted_loss = group_loss
if torch.all(self.adj>0):
adjusted_loss += self.adj/torch.sqrt(self.group_counts)
if self.normalize_loss:
adjusted_loss = adjusted_loss/(adjusted_loss.sum())
self.adv_probs = self.adv_probs * torch.exp(self.step_size*adjusted_loss.data)
self.adv_probs = self.adv_probs/(self.adv_probs.sum())
robust_loss = group_loss @ self.adv_probs
return robust_loss, self.adv_probs
def compute_robust_loss_btl(self, group_loss, group_count):
adjusted_loss = self.exp_avg_loss + self.adj/torch.sqrt(self.group_counts)
return self.compute_robust_loss_greedy(group_loss, adjusted_loss)
def compute_robust_loss_greedy(self, group_loss, ref_loss):
sorted_idx = ref_loss.sort(descending=True)[1]
sorted_loss = group_loss[sorted_idx]
sorted_frac = self.group_frac[sorted_idx]
mask = torch.cumsum(sorted_frac, dim=0)<=self.alpha
weights = mask.float() * sorted_frac /self.alpha
last_idx = mask.sum()
weights[last_idx] = 1 - weights.sum()
weights = sorted_frac*self.min_var_weight + weights*(1-self.min_var_weight)
robust_loss = sorted_loss @ weights
# sort the weights back
_, unsort_idx = sorted_idx.sort()
unsorted_weights = weights[unsort_idx]
return robust_loss, unsorted_weights
def compute_group_avg(self, losses, group_idx):
# compute observed counts and mean loss for each group
group_map = (group_idx == torch.arange(self.n_groups).unsqueeze(1).long().cuda()).float()
group_count = group_map.sum(1)
group_denom = group_count + (group_count==0).float() # avoid nans
group_loss = (group_map @ losses.view(-1))/group_denom
return group_loss, group_count
def update_exp_avg_loss(self, group_loss, group_count):
prev_weights = (1 - self.gamma*(group_count>0).float()) * (self.exp_avg_initialized>0).float()
curr_weights = 1 - prev_weights
self.exp_avg_loss = self.exp_avg_loss * prev_weights + group_loss*curr_weights
self.exp_avg_initialized = (self.exp_avg_initialized>0) + (group_count>0)
def reset_stats(self):
self.processed_data_counts = torch.zeros(self.n_groups).cuda()
self.update_data_counts = torch.zeros(self.n_groups).cuda()
self.update_batch_counts = torch.zeros(self.n_groups).cuda()
self.avg_group_loss = torch.zeros(self.n_groups).cuda()
self.avg_group_acc = torch.zeros(self.n_groups).cuda()
self.avg_per_sample_loss = 0.
self.avg_actual_loss = 0.
self.avg_acc = 0.
self.batch_count = 0.
def update_stats(self, actual_loss, group_loss, group_acc, group_count, weights=None):
# avg group loss
denom = self.processed_data_counts + group_count
denom += (denom==0).float()
prev_weight = self.processed_data_counts/denom
curr_weight = group_count/denom
self.avg_group_loss = prev_weight*self.avg_group_loss + curr_weight*group_loss
# avg group acc
self.avg_group_acc = prev_weight*self.avg_group_acc + curr_weight*group_acc
# batch-wise average actual loss
denom = self.batch_count + 1
self.avg_actual_loss = (self.batch_count/denom)*self.avg_actual_loss + (1/denom)*actual_loss
# counts
self.processed_data_counts += group_count
if self.is_robust:
self.update_data_counts += group_count*((weights>0).float())
self.update_batch_counts += ((group_count*weights)>0).float()
else:
self.update_data_counts += group_count
self.update_batch_counts += (group_count>0).float()
self.batch_count+=1
# avg per-sample quantities
group_frac = self.processed_data_counts/(self.processed_data_counts.sum())
self.avg_per_sample_loss = group_frac @ self.avg_group_loss
self.avg_acc = group_frac @ self.avg_group_acc
def get_model_stats(self, model, args, stats_dict):
model_norm_sq = 0.
for param in model.parameters():
model_norm_sq += torch.norm(param) ** 2
stats_dict['model_norm_sq'] = model_norm_sq.item()
stats_dict['reg_loss'] = args.weight_decay / 2 * model_norm_sq.item()
return stats_dict
def get_stats(self, model=None, args=None):
stats_dict = {}
for idx in range(self.n_groups):
stats_dict[f'avg_loss_group:{idx}'] = self.avg_group_loss[idx].item()
stats_dict[f'exp_avg_loss_group:{idx}'] = self.exp_avg_loss[idx].item()
stats_dict[f'avg_acc_group:{idx}'] = self.avg_group_acc[idx].item()
stats_dict[f'processed_data_count_group:{idx}'] = self.processed_data_counts[idx].item()
stats_dict[f'update_data_count_group:{idx}'] = self.update_data_counts[idx].item()
stats_dict[f'update_batch_count_group:{idx}'] = self.update_batch_counts[idx].item()
stats_dict['avg_actual_loss'] = self.avg_actual_loss.item()
stats_dict['avg_per_sample_loss'] = self.avg_per_sample_loss.item()
stats_dict['avg_acc'] = self.avg_acc.item()
# Model stats
if model is not None:
assert args is not None
stats_dict = self.get_model_stats(model, args, stats_dict)
return stats_dict
def log_stats(self, logger, is_training):
if logger is None:
return
logger.write(f'Average incurred loss: {self.avg_per_sample_loss.item():.3f} \n')
logger.write(f'Average sample loss: {self.avg_actual_loss.item():.3f} \n')
logger.write(f'Average acc: {self.avg_acc.item():.3f} \n')
for group_idx in range(self.n_groups):
logger.write(
f' {self.group_str(group_idx)} '
f'[n = {int(self.processed_data_counts[group_idx])}]:\t'
f'loss = {self.avg_group_loss[group_idx]:.3f} '
f'exp loss = {self.exp_avg_loss[group_idx]:.3f} '
f'adjusted loss = {self.exp_avg_loss[group_idx] + self.adj[group_idx]/torch.sqrt(self.group_counts)[group_idx]:.3f} '
f'adv prob = {self.adv_probs[group_idx]:3f} '
f'acc = {self.avg_group_acc[group_idx]:.3f}\n')
logger.flush()