我可以使用逻辑索引或索引列表来切片张量？

Question

我正在尝试使用列上的逻辑索引切片PyTorch张量.我想要与索引向量中的1值对应的列.切片和逻辑索引都是可能的,但它们是否可以一起使用？如果是这样,怎么样？我的尝试一直在抛出无益的错误

TypeError:使用ByteTensor类型的对象索引张量.唯一支持的类型是整数,切片,numpy标量和torch.LongTensor或torch.ByteTensor作为唯一的参数.

MCVE

期望的输出

C = torch.LongTensor([[1, 3], [4, 6]])
# 1 3
# 4 6

仅对列进行逻辑索引

import torch
A_log = torch.ByteTensor([1, 0, 1]) # the logical index
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B[:, A_log] # Throws error

我也试过使用索引列表

import torch
A_idx = torch.LongTensor([0, 2]) # the index vector
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B[:, A_idx] # Throws error

如果向量大小相同,则逻辑索引有效

import torch
A_log = torch.ByteTensor([1, 0, 1]) # the logical index
B = torch.LongTensor([1, 2, 3])
C = B[A_log]

如果我使用连续的索引范围,切片工作

import torch
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B[:, 1:2]

我可以通过重复逻辑索引来获得所需的结果,使其与我正在索引的张量具有相同的大小,但是我还必须重新整形输出.

import torch
A_log = torch.ByteTensor([1, 0, 1]) # the logical index
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B[A_log.repeat(2, 1)] # [torch.LongTensor of size 4]
C = C.resize_(2, 2)

Answer 1

我认为这是作为index_select功能实现的,你可以试试

import torch
A_idx = torch.LongTensor([0, 2]) # the index vector
B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
C = B.index_select(1, A_idx)
# 1 3
# 4 6

Answer 2

在 PyTorch 1.5.0 中，用作索引的张量必须是 long、byte 或 bool 张量。

以下是作为多头张量的索引。

import torch

B = torch.LongTensor([[1, 2, 3], [4, 5, 6]])
idx1 = torch.LongTensor([0, 2])

B[:, idx1]
# tensor([[1, 3],
#         [4, 6]])

这是一个布尔张量（逻辑索引）：

idx2 = torch.BoolTensor([True, False, True]) 

B[:, idx2]
# tensor([[1, 3],
#         [4, 6]])