让 a
成为一个 (n, d, l)
张量。让 indices
是一个 (n, 1)
张量,包含索引。我想从 a
给出的索引中的中间维度张量中收集 indices
。因此,生成的张量的形状为 (n, l)
。
n = 3
d = 2
l = 3
a = tensor([[[ 0, 1, 2],
[ 3, 4, 5]],
[[ 6, 7, 8],
[ 9, 10, 11]],
[[12, 13, 14],
[15, 16, 17]]])
indices = tensor([[0],
[1],
[0]])
# Shape of result is (n, l)
result = tensor([[ 0, 1, 2], # a[0, 0, :] since indices[0] == 0
[ 9, 10, 11], # a[1, 1, :] since indices[1] == 1
[12, 13, 14]]) # a[2, 0, :] since indices[2] == 0
这确实类似于 a.gather(1, indices)
,但 gather
不起作用,因为 indices
与 a
的形状不同。如何在此设置中使用 gather
?或者我应该用什么?
答案 0 :(得分:0)
您可以手动创建索引。如果 indices
张量具有示例数据的形状,则必须将其展平。
a[torch.arange(len(a)),indices.view(-1)]
# equal to a[[0,1,2],[0,1,0]]
出:
tensor([[ 0, 1, 2],
[ 9, 10, 11],
[12, 13, 14]])
答案 1 :(得分:0)
对于索引维度两侧的更多维度,我在 Michael 的顶部添加了我的答案,但我希望有人给我一个不使用 arange
的更好的答案!
def squeeze_index(x, dim, index):
# flatten to rows
y = x.view((-1,) + x.shape[dim:])
# generate row indices
rows = torch.arange(y.shape[0])
# index and reshape
result_shape = x.shape[:dim] + (x.shape[dim+1:] if dim != -1 else ())
return y[rows, index.view(-1), ...].view(result_shape)
a = torch.arange(2*3*2*3).reshape((2,3,2,3))
indices = torch.tensor([0,0,1,0,0,1]).reshape((2,3))
result = squeeze_index(a, 2, i)
print("a", a.shape, a)
print("indices", indices.shape, indices)
print("result", result.shape, result)
给出:
a torch.Size([2, 3, 2, 3]) tensor([[[[ 0, 1, 2],
[ 3, 4, 5]],
[[ 6, 7, 8],
[ 9, 10, 11]],
[[12, 13, 14],
[15, 16, 17]]],
[[[18, 19, 20],
[21, 22, 23]],
[[24, 25, 26],
[27, 28, 29]],
[[30, 31, 32],
[33, 34, 35]]]])
indices torch.Size([2, 3]) tensor([[0, 0, 1],
[0, 0, 1]])
result torch.Size([2, 3, 3]) tensor([[[ 0, 1, 2],
[ 6, 7, 8],
[15, 16, 17]],
[[18, 19, 20],
[24, 25, 26],
[33, 34, 35]]])
答案 2 :(得分:0)
在使用gather函数之前,重塑索引,这是一个例子
def gather_righthand(src, index, check=True):
index = index.long()
i_dim = index.dim(); s_dim = src.dim(); t_dim = i_dim-1
if check:
assert s_dim > i_dim
for d in range(0, t_dim):
assert src.shape[d] == index.shape[d]
index_new_shape = list(src.shape)
index_new_shape[t_dim] = index.shape[t_dim]
for _ in range(i_dim, s_dim): index = index.unsqueeze(-1)
index_expand = index.expand(index_new_shape) # only this two line matters
return torch.gather(src, dim=t_dim, index=index_expand) # only this two line matters
gather_righthand(a, indices)
tensor([[[ 0., 1., 2.]],
[[ 9., 10., 11.]],
[[12., 13., 14.]]])