File size: 941 Bytes
38f9214 | 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 | import torch
from torch import nn
class ClassificationHead(nn.Module):
def __init__(self, channels: int, outputs: int, bins: int):
super().__init__()
self.proj = nn.Conv2d(channels, outputs * bins, 1)
self.outputs, self.bins = outputs, bins
def forward(self, x: torch.Tensor) -> torch.Tensor:
b, _, h, w = x.shape
return self.proj(x).reshape(b, self.outputs, self.bins, h, w)
class RegressionHead(nn.Module):
def __init__(self, channels: int, outputs: int):
super().__init__()
self.proj = nn.Conv2d(channels, outputs, 1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.proj(x)
def decode_bins(logits: torch.Tensor, minimum: float = 0.0, maximum: float = 1.0) -> torch.Tensor:
centers = torch.linspace(minimum, maximum, logits.shape[2], device=logits.device)
return (logits.softmax(2) * centers[None, None, :, None, None]).sum(2)
|