| | |
| | |
| |
|
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| | import torch |
| | import numpy as np |
| |
|
| |
|
| | class GeodesicLoss(nn.Module): |
| | def __init__(self): |
| | super(GeodesicLoss, self).__init__() |
| |
|
| | def compute_geodesic_distance(self, m1, m2): |
| | """ Compute the geodesic distance between two rotation matrices. |
| | |
| | Args: |
| | m1, m2: Two rotation matrices with the shape (batch x 3 x 3). |
| | |
| | Returns: |
| | The minimal angular difference between two rotation matrices in radian form [0, pi]. |
| | """ |
| | m1 = m1.reshape(-1, 3, 3) |
| | m2 = m2.reshape(-1, 3, 3) |
| | batch = m1.shape[0] |
| | m = torch.bmm(m1, m2.transpose(1, 2)) |
| |
|
| | cos = (m[:, 0, 0] + m[:, 1, 1] + m[:, 2, 2] - 1) / 2 |
| | cos = torch.clamp(cos, min=-1 + 1E-6, max=1-1E-6) |
| |
|
| | theta = torch.acos(cos) |
| |
|
| | return theta |
| |
|
| | def __call__(self, m1, m2, reduction='mean'): |
| | loss = self.compute_geodesic_distance(m1, m2) |
| |
|
| | if reduction == 'mean': |
| | return loss.mean() |
| | elif reduction == 'none': |
| | return loss |
| | else: |
| | raise RuntimeError(f'unsupported reduction: {reduction}') |
| |
|
| |
|
| | class BCE_Loss(nn.Module): |
| | def __init__(self, args=None): |
| | super(BCE_Loss, self).__init__() |
| | |
| | def forward(self, fake_outputs, real_target): |
| | final_loss = F.cross_entropy(fake_outputs, real_target, reduce="mean") |
| | return final_loss |
| |
|
| | class weight_Loss(nn.Module): |
| | def __init__(self, args=None): |
| | super(weight_Loss, self).__init__() |
| | def forward(self, weight_f): |
| | weight_loss_div = torch.mean(weight_f[:, :, 0]*weight_f[:, :, 1]) |
| | weight_loss_gap = torch.mean(-torch.log(torch.max(weight_f[:, :, 0], dim=1)[0] - torch.min(weight_f[:, :, 0], dim=1)[0])) |
| | return weight_loss_div, weight_loss_gap |
| | |
| |
|
| | class HuberLoss(nn.Module): |
| | def __init__(self, beta=0.1, reduction="mean"): |
| | super(HuberLoss, self).__init__() |
| | self.beta = beta |
| | self.reduction = reduction |
| | |
| | def forward(self, outputs, targets): |
| | final_loss = F.smooth_l1_loss(outputs / self.beta, targets / self.beta, reduction=self.reduction) * self.beta |
| | return final_loss |
| | |
| |
|
| | class KLDLoss(nn.Module): |
| | def __init__(self, beta=0.1): |
| | super(KLDLoss, self).__init__() |
| | self.beta = beta |
| | |
| | def forward(self, outputs, targets): |
| | final_loss = F.smooth_l1_loss((outputs / self.beta, targets / self.beta) * self.beta) |
| | return final_loss |
| |
|
| |
|
| | class REGLoss(nn.Module): |
| | def __init__(self, beta=0.1): |
| | super(REGLoss, self).__init__() |
| | self.beta = beta |
| | |
| | def forward(self, outputs, targets): |
| | final_loss = F.smooth_l1_loss((outputs / self.beta, targets / self.beta) * self.beta) |
| | return final_loss |
| |
|
| |
|
| | class L2Loss(nn.Module): |
| | def __init__(self): |
| | super(L2Loss, self).__init__() |
| | |
| | def forward(self, outputs, targets): |
| | final_loss = F.l2_loss(outputs, targets) |
| | return final_loss |
| |
|
| | LOSS_FUNC_LUT = { |
| | "bce_loss": BCE_Loss, |
| | "l2_loss": L2Loss, |
| | "huber_loss": HuberLoss, |
| | "kl_loss": KLDLoss, |
| | "id_loss": REGLoss, |
| | "GeodesicLoss": GeodesicLoss, |
| | "weight_Loss": weight_Loss, |
| | } |
| |
|
| |
|
| | def get_loss_func(loss_name, **kwargs): |
| | loss_func_class = LOSS_FUNC_LUT.get(loss_name) |
| | loss_func = loss_func_class(**kwargs) |
| | return loss_func |
| |
|
| |
|
| |
|
