python - 在相等索引位置有效填充 torch.Tensor

Question

我有一个 6 维全零 pytorch 张量 lrel_w，我想在前三个维度的索引和最后三个维度的索引匹配的位置填充 1。我目前正在使用 3 个嵌套 for 循环简单地解决这个问题:

lrel_w = torch.zeros(
  input_size[0], input_size[1], input_size[2],
  input_size[0], input_size[1], input_size[2]
)
for c in range(input_size[0]):
  for x in range(input_size[1]):
    for y in range(input_size[2]):
      lrel_w[c,x,y,c,x,y] = 1

我确信一定有一种更有效的方法来做到这一点，但我还没有弄清楚。

Answer 1

你可以试试这个。

import torch
c, m, n = input_size[0], input_size[1], input_size[2]

t = torch.zeros(c, m, n, c, m, n)
i, j, k = torch.meshgrid(torch.arange(c), torch.arange(m), torch.arange(n))
i = i.flatten()
j = j.flatten()
k = k.flatten()

t[i, j, k, i, j, k] = 1

方法如下meshgrid如果您需要引用，可以使用。