python - 一个 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)。