Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import nn | |
| def generalized_mean(x, p, dim): | |
| x_type = x.dtype | |
| x = x.to(torch.double) | |
| x = x**p | |
| x = x.mean(dim=dim) | |
| x = x**(1/p) | |
| return x.to(x_type) | |
| def surject_to_positive(x, c=5): | |
| assert x.min() >= -1 | |
| assert x.max() <= 1 | |
| return c + c * x | |
| def surject_from_positive(x, c=5): | |
| return (x - c) / c | |
| class BoilerplateLoss(nn.Module): | |
| def __init__(self) -> None: | |
| super().__init__() | |
| self.p = 9 | |
| def forward(self, y_pred, y_attack, **kwargs): | |
| y_pred = y_pred.softmax(dim=-1) | |
| C = y_pred.shape[1] | |
| K = y_attack.shape[1] | |
| desired_mask = torch.zeros_like(y_pred, dtype=torch.bool) | |
| desired_mask.scatter_(dim=1, index=y_attack, | |
| src=torch.ones_like(y_attack, dtype=torch.bool)) | |
| y_not_in_attack = (~desired_mask).nonzero()[:, 1].view(-1, C - K) | |
| y_pred_in_attack = torch.gather(y_pred, dim=1, index=y_attack) | |
| y_pred_not_in_attack = torch.gather(y_pred, dim=1, index=y_not_in_attack) | |
| y_pred_in_attack_min = y_pred_in_attack.min(dim=-1).values #generalized_mean(y_pred_in_attack, -self.p, dim=1) | |
| y_pred_not_in_attack_max = y_pred_not_in_attack.max(dim=-1).values #generalized_mean(y_pred_not_in_attack, self.p, dim=1) | |
| macro_loss = (y_pred_not_in_attack_max - y_pred_in_attack_min) | |
| sorting_loss = y_pred_in_attack.diff(dim=-1) | |
| # Surject sorting_loss to positive domain, since it goes [-1,1] we can just shift by 1 | |
| sorting_loss = surject_to_positive(sorting_loss) | |
| sorting_loss = generalized_mean(sorting_loss, p=9, dim=-1) | |
| # Surject back | |
| sorting_loss = surject_from_positive(sorting_loss) | |
| catted_loss = torch.stack([macro_loss, sorting_loss], dim=-1) | |
| catted_loss_pos = surject_to_positive(catted_loss) | |
| final_loss_pos = generalized_mean(catted_loss_pos, p=10, dim=-1) | |
| final_loss = surject_from_positive(final_loss_pos) | |
| return final_loss | |