| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| try: |
| from scipy.optimize import linear_sum_assignment |
| SCIPY_AVAILABLE = True |
| except ImportError: |
| SCIPY_AVAILABLE = False |
|
|
|
|
| class TokenL2Loss(nn.Module): |
| """ |
| 输入: |
| history: [B, T, N, D]，第 0 帧到第 T-1 帧 |
| refined_pred: [B, T, N, D]，模型预测的第 1 帧到第 T 帧 |
| target: [B, T, N, D]，真实的第 1 帧到第 T 帧 |
| |
| 功能: |
| 1. 先计算 refined_pred 和 target 的原始逐 index MSE。 |
| 2. 对 history[:, t] 和 target[:, t] 的 N 个 tokens 做一一匹配。 |
| 匹配依据是 cosine similarity，目标是最大化整体相似度。 |
| 3. 根据匹配关系，将 target[:, t] 重排到 history[:, t] 的 token 顺序。 |
| 4. 计算 refined_pred 和 matched_target 之间的 MSE。 |
| """ |
|
|
| def __init__(self, matching_method="hungarian"): |
| """ |
| 参数: |
| matching_method: |
| "hungarian": 全局最优一一匹配，推荐，但依赖 scipy，速度较慢。 |
| "greedy": 按 cosine 从大到小贪心匹配，不依赖 scipy。 |
| """ |
| super().__init__() |
| if matching_method not in ["hungarian", "greedy"]: |
| raise ValueError("matching_method must be 'hungarian' or 'greedy'.") |
|
|
| if matching_method == "hungarian" and not SCIPY_AVAILABLE: |
| print("[Warning] scipy not available, fallback to greedy matching.") |
| matching_method = "greedy" |
|
|
| self.matching_method = matching_method |
|
|
| @torch.no_grad() |
| def _hungarian_match_one(self, src_tokens, tgt_tokens): |
| """ |
| src_tokens: [N, D] |
| tgt_tokens: [N, D] |
| |
| 返回: |
| matched_indices: [N] |
| matched_indices[i] = j 表示 src token i 匹配 target token j |
| """ |
| src = F.normalize(src_tokens.float(), dim=-1) |
| tgt = F.normalize(tgt_tokens.float(), dim=-1) |
|
|
| sim = torch.matmul(src, tgt.T) |
|
|
| |
| |
| cost = (-sim).detach().cpu().numpy() |
| row_ind, col_ind = linear_sum_assignment(cost) |
|
|
| matched_indices = torch.empty( |
| src_tokens.shape[0], |
| dtype=torch.long, |
| device=src_tokens.device, |
| ) |
| matched_indices[torch.from_numpy(row_ind).to(src_tokens.device)] = torch.from_numpy(col_ind).to(src_tokens.device) |
| return matched_indices |
|
|
| @torch.no_grad() |
| def _greedy_match_one(self, src_tokens, tgt_tokens): |
| """ |
| 贪心一一匹配: |
| 按 cosine similarity 从大到小排序，优先匹配最高相似度 pair。 |
| """ |
| src = F.normalize(src_tokens.float(), dim=-1) |
| tgt = F.normalize(tgt_tokens.float(), dim=-1) |
|
|
| sim = torch.matmul(src, tgt.T) |
| n = sim.shape[0] |
|
|
| flat_sim = sim.reshape(-1) |
| sorted_indices = torch.argsort(flat_sim, descending=True) |
|
|
| matched_indices = torch.full( |
| (n,), |
| -1, |
| dtype=torch.long, |
| device=src_tokens.device, |
| ) |
| used_src = torch.zeros(n, dtype=torch.bool, device=src_tokens.device) |
| used_tgt = torch.zeros(n, dtype=torch.bool, device=src_tokens.device) |
|
|
| num_matched = 0 |
| for flat_idx in sorted_indices: |
| src_idx = flat_idx // n |
| tgt_idx = flat_idx % n |
|
|
| if used_src[src_idx] or used_tgt[tgt_idx]: |
| continue |
|
|
| matched_indices[src_idx] = tgt_idx |
| used_src[src_idx] = True |
| used_tgt[tgt_idx] = True |
| num_matched += 1 |
|
|
| if num_matched == n: |
| break |
|
|
| return matched_indices |
|
|
| @torch.no_grad() |
| def build_matched_target(self, history, target): |
| """ |
| 根据 history 和 target 的相邻帧 token cosine similarity， |
| 为每个 target frame 构造按 history token 顺序重排后的 matched_target。 |
| |
| history: [B, T, N, D] |
| target: [B, T, N, D] |
| |
| 返回: |
| matched_target: [B, T, N, D] |
| match_indices: [B, T, N] |
| """ |
| if history.shape != target.shape: |
| raise ValueError( |
| f"history and target must have same shape, got " |
| f"{tuple(history.shape)} and {tuple(target.shape)}." |
| ) |
|
|
| bsz, seq_len, num_tokens, dim = history.shape |
|
|
| matched_target = torch.empty_like(target) |
| match_indices = torch.empty( |
| (bsz, seq_len, num_tokens), |
| dtype=torch.long, |
| device=target.device, |
| ) |
|
|
| for b in range(bsz): |
| for t in range(seq_len): |
| src_tokens = history[b, t] |
| tgt_tokens = target[b, t] |
|
|
| if self.matching_method == "hungarian": |
| indices = self._hungarian_match_one(src_tokens, tgt_tokens) |
| else: |
| indices = self._greedy_match_one(src_tokens, tgt_tokens) |
|
|
| match_indices[b, t] = indices |
| matched_target[b, t] = target[b, t, indices] |
|
|
| return matched_target, match_indices |
|
|
| def forward(self, history, refined_pred, target): |
| """ |
| 返回: |
| { |
| "total": matched_mse, |
| "original_mse": 原始逐 index MSE, |
| "matched_mse": 匹配后一一对应 MSE, |
| "match_indices": 匹配关系 [B, T, N] |
| } |
| """ |
| history = history.float() |
| refined_pred = refined_pred.float() |
| target = target.float() |
|
|
| if refined_pred.shape != target.shape: |
| raise ValueError( |
| f"refined_pred and target must have same shape, got " |
| f"{tuple(refined_pred.shape)} and {tuple(target.shape)}." |
| ) |
|
|
| if history.shape != target.shape: |
| raise ValueError( |
| f"history and target must have same shape, got " |
| f"{tuple(history.shape)} and {tuple(target.shape)}." |
| ) |
|
|
| original_mse = F.mse_loss(refined_pred, target) |
|
|
| matched_target, match_indices = self.build_matched_target( |
| history=history, |
| target=target, |
| ) |
|
|
| matched_mse = F.mse_loss(refined_pred, matched_target) |
|
|
| return { |
| "total": matched_mse, |
| "original_mse": original_mse, |
| "matched_mse": matched_mse, |
| "match_indices": match_indices, |
| } |
