我有一个索引列表和一个形状为:
shape = [batch_size, d_0, d_1, ..., d_k]
idx = [i_0, i_1, ..., i_k]
有没有一种有效的方法,可以在每个维度 d_0, ..., d_k
上使用索引 i_0, ..., i_k
进行高效的张量索引?(k
只能在运行时确定)
结果应该是:
tensor[:, i_0, i_1, ..., i_k] #tensor.shape = [batch_size]
目前我正在创建一个切片的元组,每个维度都有一个切片:
idx = (slice(tensor.shape[0]),) + tuple(slice(i, i+1) for i in idx)
tensor[idx]
但我更喜欢像这样的东西:
tensor[:, *idx]
抱歉,我只能使用英语回答问题。
a = torch.randint(0,10,[3,3,3,3])
indexes = torch.LongTensor([1,1,1])
我想仅索引最后 len(indexes) 维度,如下所示:
a[:, indexes[0], indexes[1], indexes[2]]
但是在一般情况下,我不知道 indexes
有多长。
注意:这个答案并不适用于所有维度,它索引了所有的维度,而且无法处理一个合适的子集!