时间戳

首页 > 解决方案 > Pytorch - 索引一系列多个索引？

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

问题描述

假设我有一个大小为 [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 循环，这实际上可能吗？

标签： pytorch

解决方案

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

以下是一些可能有用的替代方案，具体取决于您的用例。特别是如果您希望重复使用相同的任务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

推荐阅读