File size: 5,965 Bytes
6979012 | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | import torch
import torch.nn.functional as F
def convert_keypoints_to_(
gt: torch.Tensor, H: int = 768, W: int = 1024
) -> torch.Tensor:
B, N, _ = gt.shape
xy = gt.round().long()
x = xy[..., 0].clamp(0, W - 1)
y = xy[..., 1].clamp(0, H - 1)
keypoint_map = torch.zeros((B, H, W), device=gt.device, dtype=torch.float32)
b = torch.arange(B, device=gt.device).unsqueeze(1).expand(B, N)
keypoint_map[b, y, x] = 1.0
return keypoint_map
def detector_loss(
keypoint_map: torch.Tensor, # (B,H,W) or (B,1,H,W), binary/bool
logits: torch.Tensor, # (B,65,Hc,Wc), raw convPb output
valid_mask: torch.Tensor | None = None, # (B,H,W) or (B,1,H,W)
grid_size: int = 8,
eps: float = 1e-8,
) -> torch.Tensor:
# --- labels: space_to_depth + dustbin + random tie-break + argmax ---
if keypoint_map.ndim == 3:
keypoint_map = keypoint_map.unsqueeze(1) # (B,1,H,W)
keypoint_map = keypoint_map.float()
# TF NHWC space_to_depth -> PyTorch NCHW pixel_unshuffle
labels = F.pixel_unshuffle(keypoint_map, downscale_factor=grid_size) # (B,64,Hc,Wc)
dustbin = torch.ones_like(labels[:, :1]) # (B,1,Hc,Wc)
labels = torch.cat([2.0 * labels, dustbin], dim=1) # (B,65,Hc,Wc)
# same tie-break idea as TF random_uniform(..., 0, 0.1)
labels = torch.argmax(
labels + 0.1 * torch.rand_like(labels), dim=1
) # (B,Hc,Wc), long
# --- valid mask path ---
if valid_mask is None:
valid_mask = torch.ones_like(keypoint_map)
elif valid_mask.ndim == 3:
valid_mask = valid_mask.unsqueeze(1)
valid_mask = valid_mask.float()
valid_mask = F.pixel_unshuffle(
valid_mask, downscale_factor=grid_size
) # (B,64,Hc,Wc)
valid_mask = torch.prod(valid_mask, dim=1) # (B,Hc,Wc)
# --- sparse softmax cross entropy with weights ---
per_cell = F.cross_entropy(logits, labels, reduction="none") # (B,Hc,Wc)
weighted = per_cell * valid_mask
loss = weighted.sum() / valid_mask.sum().clamp_min(1.0 + eps)
return loss
def detector_loss_simple(
gt_logits: torch.Tensor,
pred_logits: torch.Tensor, # (B,65,Hc,Wc), raw convPb output
) -> torch.Tensor:
gt_labels = torch.argmax(
gt_logits + 0.1 * torch.rand_like(gt_logits), dim=1
) # (B,Hc,Wc), long
return F.cross_entropy(pred_logits, gt_labels, reduction="mean")
def detector_kd_kl(teacher_logits, student_logits, T=2.0, valid_mask=None, eps=1e-8):
# student_logits, teacher_logits: (B,65,Hc,Wc)
log_p_s = F.log_softmax(student_logits / T, dim=1)
p_t = F.softmax(teacher_logits / T, dim=1)
# KL per cell: (B,Hc,Wc)
kl_map = F.kl_div(log_p_s, p_t, reduction="none").sum(dim=1)
if valid_mask is None:
return (T * T) * kl_map.mean()
# valid_mask expected (B,H,W) or (B,1,H,W), convert to cell mask (B,Hc,Wc)
if valid_mask.ndim == 3:
valid_mask = valid_mask.unsqueeze(1)
vm = F.pixel_unshuffle(valid_mask.float(), downscale_factor=8) # (B,64,Hc,Wc)
vm = torch.prod(vm, dim=1) # (B,Hc,Wc)
return (T * T) * (kl_map * vm).sum() / vm.sum().clamp_min(1.0 + eps)
def descriptor_loss_simple(
pred_desc: torch.Tensor, # (B, D, Hc, Wc)
gt_desc: torch.Tensor, # (B, D, Hc, Wc)
valid_mask: torch.Tensor | None = None, # (B,H,W) or (B,1,H,W)
grid_size: int = 8,
eps: float = 1e-8,
) -> torch.Tensor:
pred = F.normalize(pred_desc, p=2, dim=1)
gt = F.normalize(gt_desc, p=2, dim=1)
# cosine distance per cell
per_cell = 1.0 - (pred * gt).sum(dim=1) # (B, Hc, Wc)
if valid_mask is None:
return per_cell.mean()
if valid_mask.ndim == 3:
valid_mask = valid_mask.unsqueeze(1) # (B,1,H,W)
vm = F.pixel_unshuffle(
valid_mask.float(), downscale_factor=grid_size
) # (B,64,Hc,Wc)
vm = torch.prod(vm, dim=1) # (B,Hc,Wc)
return (per_cell * vm).sum() / vm.sum().clamp_min(1.0 + eps)
def descriptor_loss(
descriptors: torch.Tensor, # (B, D, Hc, Wc), student
target_descriptors: torch.Tensor, # (B, D, Hc, Wc), teacher/GT
valid_mask: torch.Tensor | None = None, # (B,H,W) or (B,1,H,W)
grid_size: int = 8,
positive_margin: float = 1.0,
negative_margin: float = 0.2,
lambda_d: float = 0.05,
eps: float = 1e-8,
) -> torch.Tensor:
B, D, Hc, Wc = descriptors.shape
HW = Hc * Wc
# L2 normalize descriptors
desc = F.normalize(descriptors, p=2, dim=1) # (B,D,Hc,Wc)
tgt = F.normalize(target_descriptors, p=2, dim=1) # (B,D,Hc,Wc)
# Flatten spatial dims
desc = desc.flatten(2).transpose(1, 2) # (B,HW,D)
tgt = tgt.flatten(2).transpose(1, 2) # (B,HW,D)
# Pairwise dot products: (B,HW,HW)
dot = torch.bmm(desc, tgt.transpose(1, 2))
dot = F.relu(dot)
# TF does double normalization over pairwise axes
dot = F.normalize(dot, p=2, dim=2)
dot = F.normalize(dot, p=2, dim=1)
# Identity correspondence mask s (diagonal)
eye = torch.eye(HW, device=dot.device, dtype=dot.dtype).unsqueeze(0) # (1,HW,HW)
s = eye.expand(B, -1, -1)
positive_dist = F.relu(positive_margin - dot)
negative_dist = F.relu(dot - negative_margin)
pairwise_loss = (
lambda_d * s * positive_dist + (1.0 - s) * negative_dist
) # (B,HW,HW)
# valid mask: same logic as TF space_to_depth + reduce_prod
if valid_mask is None:
vm = torch.ones(
(B, Hc * grid_size, Wc * grid_size), device=dot.device, dtype=dot.dtype
)
else:
vm = valid_mask
if vm.ndim == 3:
vm = vm.unsqueeze(1) # (B,1,H,W)
vm = vm.float()
vm = F.pixel_unshuffle(vm, downscale_factor=grid_size) # (B,grid^2,Hc,Wc)
vm = torch.prod(vm, dim=1) # (B,Hc,Wc)
vm = vm.reshape(B, HW) # valid target cells
vm = vm[:, None, :] # (B,1,HW), broadcast to (B,HW,HW)
normalization = vm.sum() * float(HW) + eps
loss = (pairwise_loss * vm).sum() / normalization
return loss
|