无需遍历向量的复杂乘法实现：PyTorch自定义多分类损失高效计算

无遍历高效实现方案 #

完全对齐你现有的遍历计算逻辑，且适配任意样本数N，不会触发维度不匹配报错的向量化实现代码如下：

# 批量取出每个样本标签对应的类别距离向量，输出形状为[N, 类别数]
sample_distance = d_mtx[y]
# 逐元素相乘后按行求和（等价于逐样本做点积），累加得到总损失
loss = (sample_distance * y_pred).sum()

逻辑说明

你之前遇到的维度报错是因为旧方案没有先按标签做索引对齐样本维度，上述实现通过d_mtx[y]直接批量拉取所有样本对应类别的距离向量，得到的张量和y_pred形状完全一致，不管样本数是3还是4或者其他数值都可以正常运行，计算结果和你写的遍历逻辑完全等价。

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