如何在PyTorch中实现自定义平均定位误差损失函数？

你需要的损失本质是预测坐标与真实坐标的欧氏距离在批量维度的平均值，可通过以下几种方式实现，直接替换原有MSE损失即可：

最简函数式实现 #

import torch

def localisation_loss(y_pred, y_true):
    # 输入参数说明：
    # y_pred 维度为 [batch_size, 2]，每行对应一个样本的预测x、y坐标
    # y_true 维度为 [batch_size, 2]，每行对应一个样本的真实x、y坐标
    squared_diff = (y_pred - y_true) ** 2
    # 按坐标维度求和后开方，得到每个样本的定位误差距离
    distance_error = torch.sqrt(torch.sum(squared_diff, dim=1))
    # 对批量内所有样本的误差取平均得到最终损失
    return torch.mean(distance_error)

训练阶段调用示例：

optimizer.zero_grad()
pred_coords = model(wifi_rss_input)
loss = localisation_loss(pred_coords, real_coords)
loss.backward()
optimizer.step()

数值稳定优化版本 #

当预测坐标与真实坐标完全相等时，torch.sqrt的输入为0，反向传播时会出现梯度NaN的问题，可通过添加极小常量epsilon解决：

def localisation_loss(y_pred, y_true, eps=1e-6):
    squared_diff = (y_pred - y_true) ** 2
    sum_squared = torch.sum(squared_diff, dim=1)
    distance_error = torch.sqrt(sum_squared + eps)
    return torch.mean(distance_error)

封装为PyTorch损失类 #

如果希望和内置损失函数用法保持一致，可封装为nn.Module的子类：

import torch.nn as nn

class LocalisationLoss(nn.Module):
    def __init__(self, eps=1e-6):
        super().__init__()
        self.eps = eps
    
    def forward(self, y_pred, y_true):
        squared_diff = (y_pred - y_true) ** 2
        sum_squared = torch.sum(squared_diff, dim=1)
        distance_error = torch.sqrt(sum_squared + self.eps)
        return torch.mean(distance_error)

调用方式和MSE完全一致：

criterion = LocalisationLoss()
loss = criterion(pred_coords, real_coords)

内容的提问来源于stack exchange，提问作者Abhiram C D