|
"""up decoder head. |
|
|
|
Adapted from https://github.com/jinlinyi/PerspectiveFields |
|
""" |
|
|
|
import logging |
|
|
|
import torch |
|
from torch import nn |
|
from torch.nn import functional as F |
|
|
|
from siclib.models import get_model |
|
from siclib.models.base_model import BaseModel |
|
from siclib.models.utils.metrics import up_error |
|
from siclib.models.utils.perspective_encoding import decode_up_bin |
|
from siclib.utils.conversions import deg2rad |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
|
|
|
|
|
|
class UpDecoder(BaseModel): |
|
default_conf = { |
|
"loss_type": "l1", |
|
"use_loss": True, |
|
"use_uncertainty_loss": True, |
|
"loss_weight": 1.0, |
|
"recall_thresholds": [1, 3, 5, 10], |
|
"decoder": {"name": "decoders.light_hamburger", "predict_uncertainty": True}, |
|
} |
|
|
|
required_data_keys = ["features"] |
|
|
|
def _init(self, conf): |
|
self.loss_type = conf.loss_type |
|
self.loss_weight = conf.loss_weight |
|
|
|
self.use_uncertainty_loss = conf.use_uncertainty_loss |
|
self.predict_uncertainty = conf.decoder.predict_uncertainty |
|
|
|
self.num_classes = 2 |
|
self.is_classification = self.conf.loss_type == "classification" |
|
if self.is_classification: |
|
self.num_classes = 73 |
|
|
|
self.decoder = get_model(conf.decoder.name)(conf.decoder) |
|
self.linear_pred_up = nn.Conv2d(self.decoder.out_channels, self.num_classes, kernel_size=1) |
|
|
|
def calculate_losses(self, predictions, targets, confidence=None): |
|
predictions = predictions.float() |
|
|
|
residuals = predictions - targets |
|
if self.loss_type == "l2": |
|
loss = (residuals**2).sum(axis=1) |
|
elif self.loss_type == "l1": |
|
loss = residuals.abs().sum(axis=1) |
|
elif self.loss_type == "dot": |
|
loss = 1 - (residuals * targets).sum(axis=1) |
|
elif self.loss_type == "cauchy": |
|
c = 0.007 |
|
residuals = (residuals**2).sum(axis=1) |
|
loss = c**2 / 2 * torch.log(1 + residuals / c**2) |
|
elif self.loss_type == "huber": |
|
c = deg2rad(1) |
|
loss = nn.HuberLoss(reduction="none", delta=c)(predictions, targets).sum(axis=1) |
|
else: |
|
raise NotImplementedError(f"Unknown loss type {self.conf.loss_type}") |
|
|
|
if confidence is not None and self.use_uncertainty_loss: |
|
conf_weight = confidence / confidence.sum(axis=(-2, -1), keepdims=True) |
|
conf_weight = conf_weight * (conf_weight.size(-1) * conf_weight.size(-2)) |
|
loss = loss * conf_weight.detach() |
|
|
|
losses = {f"up-{self.loss_type}-loss": loss.mean(axis=(1, 2))} |
|
losses = {k: v * self.loss_weight for k, v in losses.items()} |
|
|
|
return losses |
|
|
|
def _forward(self, data): |
|
out = {} |
|
x, log_confidence = self.decoder(data["features"]) |
|
up = self.linear_pred_up(x) |
|
|
|
if self.predict_uncertainty: |
|
out["up_confidence"] = torch.sigmoid(log_confidence) |
|
|
|
if self.is_classification: |
|
out["up_field"] = decode_up_bin(up.argmax(dim=1), self.num_classes) |
|
return out |
|
|
|
up = F.normalize(up, dim=1) |
|
|
|
out["up_field"] = up |
|
return out |
|
|
|
def loss(self, pred, data): |
|
if not self.conf.use_loss or self.is_classification: |
|
return {}, self.metrics(pred, data) |
|
|
|
predictions = pred["up_field"] |
|
targets = data["up_field"] |
|
|
|
losses = self.calculate_losses(predictions, targets, pred.get("up_confidence")) |
|
|
|
total = 0 + losses[f"up-{self.loss_type}-loss"] |
|
losses |= {"up_total": total} |
|
return losses, self.metrics(pred, data) |
|
|
|
def metrics(self, pred, data): |
|
predictions = pred["up_field"] |
|
targets = data["up_field"] |
|
|
|
mask = predictions.sum(axis=1) != 0 |
|
|
|
error = up_error(predictions, targets) * mask |
|
out = {"up_angle_error": error.mean(axis=(1, 2))} |
|
|
|
if "up_confidence" in pred: |
|
weighted_error = (error * pred["up_confidence"]).sum(axis=(1, 2)) |
|
out["up_angle_error_weighted"] = weighted_error / pred["up_confidence"].sum(axis=(1, 2)) |
|
|
|
for th in self.conf.recall_thresholds: |
|
rec = (error < th).float().mean(axis=(1, 2)) |
|
out[f"up_angle_recall@{th}"] = rec |
|
|
|
return out |
|
|
