[RFC] Fused RoIAlign Face Embedding Extractor & Numerically Stable Angular Margin Loss (`torchvision.ops` & `torchvision.losses`)

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.ops import roi_align

class AdaptiveArcFaceLoss(nn.Module):
    """
    Numerically stable ArcFace Loss enforcing safe-clamping bounds on arccos inputs
    to eliminate NaN gradient explosions in FP16/FP32 training.
    """
    def __init__(self, in_features: int, num_classes: int, scale: float = 64.0, margin: float = 0.50):
        super().__init__()
        self.scale = scale
        self.cos_m, self.sin_m = math.cos(margin), math.sin(margin)
        self.th = math.cos(math.pi - margin)
        self.mm = math.sin(math.pi - margin) * margin
        self.weight = nn.Parameter(torch.empty(num_classes, in_features))
        nn.init.xavier_uniform_(self.weight)

    def forward(self, embeddings: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        # EPS Safeguard: Prevent float overflow outside [-1.0, 1.0] bounds
        cosine = F.linear(F.normalize(embeddings), F.normalize(self.weight)).clamp(-1.0 + 1e-7, 1.0 - 1e-7)
        sine = torch.sqrt((1.0 - torch.pow(cosine, 2)).clamp(0.0, 1.0))
        phi = torch.where(cosine > self.th, cosine * self.cos_m - sine * self.sin_m, cosine - self.mm)
        
        one_hot = torch.zeros_like(cosine).scatter_(1, labels.view(-1, 1).long(), 1.0)
        logits = (one_hot * phi) + ((1.0 - one_hot) * cosine)
        return F.cross_entropy(logits * self.scale, labels)


class VectorizedFaceExtractor(nn.Module):
    """
    Fused GPU RoI crop & normalize wrapper operating zero-copy in VRAM.
    """
    def __init__(self, backbone: nn.Module, target_size=(112, 112)):
        super().__init__()
        self.backbone = backbone
        self.target_size = target_size

    def forward(self, images: torch.Tensor, boxes: torch.Tensor, box_indices: torch.Tensor) -> torch.Tensor:
        rois = torch.cat([box_indices.unsqueeze(1).float(), boxes], dim=1)
        aligned = roi_align(images, rois, output_size=self.target_size, spatial_scale=1.0)
        return F.normalize(self.backbone(aligned), p=2, dim=1)