Pytorch - 索引一系列多个索引？

Question

假设我有一个大小为 [100, 100] 的张量，并且我有一组大小为 [100] 的start_indices和。end_indices

我希望能够做这样的事情：

tensor[start_indices:end_indices, :] = 0

不幸的是，我收到一条错误消息

TypeError: only integer tensors of a single element can be converted to an index

那么，如果没有 for 循环，这实际上可能吗？

Answer 1

没有任何类型的循环

mytensor = torch.randn(4, 10)
start = torch.tensor([1, 3, 1, 3]).unsqueeze(-1)
end = torch.tensor([3, 5, 7, 4]).unsqueeze(-1)

创建与要索引的张量形状相同的索引数组

index = torch.arange(10).repeat(4).reshape(4, 10)
>>> index

tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])

然后简单地转换为布尔掩码

gte_start = start <= index
lt_end = index < end
mask = gte_start & lt_end
>>> mask
tensor([[False,  True,  True,  True, False, False, False, False, False, False],
        [False, False, False,  True,  True,  True, False, False, False, False],
        [False,  True,  True,  True,  True,  True,  True,  True, False, False],
        [False, False, False,  True,  True, False, False, False, False, False]])

并应用它

mytensor[mask] = 42.0

如果不需要循环（通常是这样），掩码是我进行多维操作的首选

Answer 2

据我所知，如果没有某种循环或列表理解，这是不可能的。

以下是一些可能有用的替代方案，具体取决于您的用例。具体来说，如果您希望重复使用相同的任务start_indices并end_indices进行多项分配，或者您希望只有一项就地分配，那么tensor下面的解决方案将很有用。

例如，如果给你一个索引列表，而start_indices不是end_indices

row_indices = torch.cat([torch.arange(s, e, dtype=torch.int64) for s, e in zip(start_indices, end_indices)])

那么这将可以使用

tensor[row_indices, :] = 0

或者如果给你一个面具

mask = torch.zeros(tensor.shape, dtype=torch.bool, device=tensor.device)
for s, e in zip(start_indices, end_indices):
    mask[s:e, :] = True

那么这将可以使用

tensor[mask] = 0