[英]PyTorch: index 2D tensor with 2D tensor of row indices
我有一个形状(x, n)
的火炬张量a
和另一个形状(y, n)
张量b
,其中y <= x
。 的每一列b
包含行索引的序列对a
和我想能够做的是某种方式索引a
与b
,使得我获得形状的张量(y, n)
其中第i列包含a[:, i][b[:, i]]
(不太确定这是否是表达它的正确方法)。
这是一个例子(其中x
= 5, y
= 3和n
= 4):
import torch
a = torch.Tensor(
[[0.1, 0.2, 0.3, 0.4],
[0.6, 0.7, 0.8, 0.9],
[1.1, 1.2, 1.3, 1.4],
[1.6, 1.7, 1.8, 1.9],
[2.1, 2.2, 2.3, 2.4]]
)
b = torch.LongTensor(
[[0, 3, 1, 2],
[2, 2, 2, 0],
[1, 1, 0, 4]]
)
# How do I get from a and b to c
# (so that I can also assign to those elements in a)?
c = torch.Tensor(
[[0.1, 1.7, 0.8, 1.4],
[1.1, 1.2, 1.3, 0.4],
[0.6, 0.7, 0.3, 2.4]]
)
我无法理解这一点。 我正在寻找的是不会产生张量的方法c
,但也让我分配相同形状的张量c
到的元素a
它c
是由。
我尝试使用index_select
但它只支持索引的1-dim
数组。
bt = b.transpose(0, 1)
at = a.transpose(0, 1)
ct = [torch.index_select(at[i], dim=0, index=bt[i]) for i in range(len(at))]
c = torch.stack(ct).transpose(0, 1)
print(c)
"""
tensor([[0.1000, 1.7000, 0.8000, 1.4000],
[1.1000, 1.2000, 1.3000, 0.4000],
[0.6000, 0.7000, 0.3000, 2.4000]])
"""
它可能不是最好的解决方案,但希望这至少可以帮到你。
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.