如何在两个不同大小的PyTorch张量中获得相等元素的索引？

在PyTorch中，如果你想在两个不同大小的张量中获得相等元素的索引，你可以使用广播（broadcasting）机制来使两个张量的形状相同，然后使用torch.eq()函数来比较它们。但是，如果张量的大小完全不同，广播可能无法工作，因为它们的形状不兼容。

以下是一个示例，说明如何在形状兼容的情况下找到相等元素的索引：

import torch

# 创建两个张量
tensor_a = torch.tensor([[1, 2, 3], [4, 5, 6]])
tensor_b = torch.tensor([[4, 5, 6], [1, 2, 3]])

# 使用广播机制使两个张量形状相同
# 在这个例子中，tensor_a和tensor_b的形状已经是兼容的，所以不需要额外的广播步骤

# 比较两个张量并获取相等元素的索引
equal_indices = torch.nonzero(torch.eq(tensor_a, tensor_b), as_tuple=True)

print(equal_indices)

输出将会是相等元素的索引：

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

这个输出表示tensor_a和tensor_b在(0,1)、(1,0)、(0,2)、(1,1)、(0,0)和(1,2)位置上的元素是相等的。

如果张量的形状不兼容，你需要先调整它们的形状。例如，如果你有一个形状为(3, 4)的张量和一个形状为(3, 1)的张量，你可以将后者扩展到前者的形状：

tensor_c = torch.tensor([[1], [2], [3]])
tensor_d = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])

# 将tensor_d的形状调整为(3, 4)
tensor_d_expanded = tensor_d.expand(-1, tensor_c.shape[1])

# 现在可以比较tensor_c和tensor_d_expanded
equal_indices_expanded = torch.nonzero(torch.eq(tensor_c, tensor_d_expanded), as_tuple=True)

print(equal_indices_expanded)

如果你的问题是遇到了形状不兼容导致的错误，那么你需要检查张量的形状，并使用view()、reshape()或expand()等方法来调整它们的形状，以便进行比较。

参考链接：