一个 3-d 张量如何被两个 2d 张量索引？

Question

这是 DIM

中第 15-20 行的快照

def random_permute(X):
    X = X.transpose(1, 2)
    b = torch.rand((X.size(0), X.size(1))).cuda()
    idx = b.sort(0)[1]
    adx = torch.range(0, X.size(1) - 1).long()
    X = X[idx, adx[None, :]].transpose(1, 2)

    return X

其中 X 是大小为 [64, 64, 128] 的张量，idx 大小为 [64, 64] 的张量，adx 大小为 [64 ]. X = X[idx, adx[None, :]] 是如何工作的？我们如何使用两个 2d 张量来索引 3d 张量？在这个索引之后 X 到底发生了什么？

Answer 1

我猜X一定是3D张量，因为它通常代表一批训练数据。

就此函数的功能而言，它随机置换输入数据张量 X 并使用以下步骤执行此操作：

首先，它使用从均匀分布中采样的值初始化张量 b。
接下来这个张量沿维度 0 排序，排序索引被拉出到张量 idx。
张量adx只是一个整数张量，取值范围从0到63。

现在，下面一行是所有魔法发生的地方：

X[idx, adx[None, :]].transpose(1, 2)

我们使用之前得到的索引idx和adx（adx[None, :]只是一个二维的行向量）。一旦我们有了它，我们就可以完全像我们在函数开头所做的那样调换轴 1 和 2：

X = X.transpose(1, 2)

这里是一个人为的例子，以便更好地理解：

# our input tensor
In [51]: X = torch.rand(64, 64, 32)

In [52]: X = X.transpose(1, 2)

In [53]: X.shape
Out[53]: torch.Size([64, 32, 64])

In [54]: b = torch.rand((X.size(0), X.size(1)))

# sort `b` which returns a tuple and take only indices
In [55]: idx = b.sort(0)[1]

In [56]: idx.shape
Out[56]: torch.Size([64, 32])

In [57]: adx = torch.arange(0, X.size(1)).long()

In [58]: adx.shape
Out[58]: torch.Size([32])

In [59]: X[idx, adx[None, :]].transpose(1, 2).shape
Out[59]: torch.Size([64, 64, 32])

这里要注意的重要一点是我们如何在最后一步中获得与输入张量形状相同的形状，即 (64, 64, 32).

Answer 2

如果我们考虑一个更小的具体例子，事情会更清楚。让

x = np.arange(8).reshape(2, 2, 2)
b = np.random.rand(2, 2)
idx = b.argsort(0) # e.g. idx=[[1, 1], [0, 0]]
adx = np.arange(2)[None, :] # [[0, 1]]
y = x[idx, adx] # implicitly expanding 'adx' to [[0, 1], [0, 1]]

在这个例子中，我们将 y 作为

y[0, 0] = x[idx[0, 0], adx[0, 0]]=x[1, 0]
y[0, 1] = x[idx[0, 1], adx[0, 1]]=x[1, 1]
y[1, 0] = x[idx[1, 0], adx[1, 0]]=x[0, 0]
...

看看我们如何在 tensorflow 中做同样的事情可能会有所帮助：

d0, d1, d2 = x.shape.as_list()
b = np.random.rand(d0, d1)
idx = np.argsort(b, 0)
idx = idx.reshape(-1)
adx = np.arange(0, d1)
adx = np.tile(adx, d0)

y = tf.reshape(tf.gather_nd(x, zip(idx, adx)), (d0, d1, d2))

一个 3-d 张量如何被两个 2d 张量索引？

How is a 3-d tensor indexed by two 2d tensors?

python

arrays

indexing

pytorch

tensor