| from typing import Any |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .base_postprocessor import BasePostprocessor |
|
|
|
|
| def kl_div(d1, d2): |
| """Compute KL-Divergence between d1 and d2.""" |
| dirty_logs = d1 * torch.log2(d1 / d2) |
| return torch.sum(torch.where(d1 != 0, dirty_logs, torch.zeros_like(d1)), |
| axis=1) |
|
|
|
|
| class RotPredPostprocessor(BasePostprocessor): |
| def __init__(self, config): |
| super(RotPredPostprocessor, self).__init__(config) |
| self.config = config |
|
|
| @torch.no_grad() |
| def postprocess(self, net: nn.Module, data: Any): |
| batch_size = len(data) |
|
|
| x_90 = torch.rot90(data, 1, [2, 3]) |
| x_180 = torch.rot90(data, 2, [2, 3]) |
| x_270 = torch.rot90(data, 3, [2, 3]) |
|
|
| x_rot = torch.cat([data, x_90, x_180, x_270]) |
| y_rot = torch.cat([ |
| torch.zeros(batch_size), |
| torch.ones(batch_size), |
| 2 * torch.ones(batch_size), |
| 3 * torch.ones(batch_size), |
| ]).long().cuda() |
|
|
| logits, logits_rot = net(x_rot, return_rot_logits=True) |
| logits = logits[:batch_size] |
| preds = logits.argmax(1) |
|
|
| |
| num_classes = logits.shape[1] |
| uniform_dist = torch.ones_like(logits) / num_classes |
| cls_loss = kl_div(uniform_dist, F.softmax(logits, dim=1)) |
|
|
| rot_one_hot = torch.zeros_like(logits_rot).scatter_( |
| 1, |
| y_rot.unsqueeze(1).cuda(), 1) |
| rot_loss = kl_div(rot_one_hot, F.softmax(logits_rot, dim=1)) |
| rot_0_loss, rot_90_loss, rot_180_loss, rot_270_loss = torch.chunk( |
| rot_loss, 4, dim=0) |
| total_rot_loss = (rot_0_loss + rot_90_loss + rot_180_loss + |
| rot_270_loss) / 4.0 |
|
|
| |
| scores = cls_loss - total_rot_loss |
| return preds, scores |
|
|
