PyTorch：基于条件实现不同形状张量的广播转换

针对将N×C的Tensor转换为N×N矩阵的需求，这里提供基于PyTorch广播机制的实现，无需循环，可高效处理N=32、64等较大规模的输入：

核心思路 #

我们需要判断任意两个行向量labels[i]和labels[j]是否存在至少一个对应位置的值相同（同为0或同为1）：

1. 利用维度扩展将原Tensor分别变为N×1×C和1×N×C的形状，触发广播机制实现所有行对的逐元素比较；
2. 对逐元素相等的结果沿C维度求和，只要总和大于0则说明存在相同位置，对应结果设为1，否则为0。

代码实现 #

import torch

# 输入的N×C Tensor
labels = torch.tensor([[1,1],[0,1],[0,0]])

# 广播实现转换
output = (labels.unsqueeze(1) == labels.unsqueeze(0)).sum(dim=-1).gt(0).float()

print(output)

输出结果 #

tensor([[1., 1., 0.],
        [1., 1., 1.],
        [0., 1., 1.]])

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