PyTorch实现Wasserstein损失：如何沿轴向量化消除批次for循环

直接替换forward方法即可，完整无循环的实现如下：

class WassersteinLoss(nn.Module):
    def __init__(self):
        super(WassersteinLoss, self).__init__()
    
    def forward(self, u, v):
        # 批量排序每个样本的分布
        u_sort_idx = torch.argsort(u, dim=1)
        v_sort_idx = torch.argsort(v, dim=1)
        u_sorted = torch.gather(u, 1, u_sort_idx)
        v_sorted = torch.gather(v, 1, v_sort_idx)
        
        # 批量拼接两个分布并排序
        all_values = torch.cat([u_sorted, v_sorted], dim=1)
        sorted_all = torch.sort(all_values, dim=1).values
        
        # 批量计算相邻值的差值
        deltas = sorted_all[:, 1:] - sorted_all[:, :-1]
        
        # 批量计算CDF索引，searchsorted自动按批次维度独立执行搜索
        u_cdf_idx = torch.searchsorted(u_sorted, sorted_all[:, :-1], right=True)
        v_cdf_idx = torch.searchsorted(v_sorted, sorted_all[:, :-1], right=True)
        
        # 批量计算CDF
        n_points = u.shape[1]
        u_cdf = u_cdf_idx / n_points
        v_cdf = v_cdf_idx / n_points
        
        # 批量计算每个样本的损失，最后取批次均值
        per_sample_loss = (torch.abs(u_cdf - v_cdf) * deltas).sum(dim=1)
        return per_sample_loss.mean()

实现对应关系（和原循环逻辑完全等价） #

• 原逐样本argsort：替换为沿dim=1的批量排序，配合gather直接取到所有样本排序后的分布值，不需要逐样本索引取值。
• 原逐样本拼接、排序u、v分布：替换为沿dim=1的批量拼接、批量排序，一次性处理所有批次样本的合并值排序。
• 原逐样本相邻值差计算：直接沿dim=1做切片差减，一次得到所有样本的deltas张量。
• 原核心的逐样本searchsorted：PyTorch内置的searchsorted原生支持批次对齐——当待搜索的基准分布（u_sorted/v_sorted）和待插入值的批次维度完全一致时，会自动对每个批次样本独立执行搜索逻辑，不需要任何Python层循环。
• 原逐样本CDF计算、逐样本损失求和：全部替换为沿dim=1的逐元素运算和维度求和，最终的均值计算和原逻辑完全一致。

验证 #

使用给出的测试用例运行修改后的代码，输出依然为tensor(5.)，和Scipy的计算结果、原循环版本的结果完全一致。全量批量运算无Python层循环，在GPU上运行时，大批次下的速度会比原循环版本提升一个数量级以上。

内容的提问来源于stack exchange，提问作者sineman