如何在火炬张量中选择特定行

Question

我想在一个 3 维的 Torch Tensor 中选择特定的行，如下所示： 它可以从每个子二维张量中获取随机数量的行，我想将它们合并在一起。 我想要的是一个 3 维张量：

tensor([[[0.7185, 0.2953, 0.6841, 0.1045]],

        [[0.6817, 0.4053, 0.2318, 0.1309]],

        [[0.8265, 0.8029, 0.3165, 0.2020],
        [0.6041, 0.0118, 0.8386, 0.5076]],

        [[0.6985, 0.7313, 0.4613, 0.1862],
        [0.5678, 0.4485, 0.4514, 0.7747]]])

但这是二维结果：

tensor([[0.7185, 0.2953, 0.6841, 0.1045],
        [0.6817, 0.4053, 0.2318, 0.1309],
        [0.8265, 0.8029, 0.3165, 0.2020],
        [0.6041, 0.0118, 0.8386, 0.5076],
        [0.6985, 0.7313, 0.4613, 0.1862],
        [0.5678, 0.4485, 0.4514, 0.7747]])

演示代码：

import torch
a= torch.rand(4,4,4)
a
tensor([[[0.1227, 0.8073, 0.0308, 0.6210],
         [0.7185, 0.2953, 0.6841, 0.1045],
         [0.2089, 0.3731, 0.7066, 0.9211],
         [0.0326, 0.4471, 0.8805, 0.3516]],     1
--------------------------------------------------
        [[0.3817, 0.2368, 0.9351, 0.0448],
         [0.6817, 0.4053, 0.2318, 0.1309],
         [0.1490, 0.4178, 0.2769, 0.7073],
         [0.0593, 0.6327, 0.0792, 0.3341]],     2
--------------------------------------------------
        [[0.3492, 0.0924, 0.8318, 0.4404],
         [0.8265, 0.8029, 0.3165, 0.2020],
         [0.6041, 0.0118, 0.8386, 0.5076],
         [0.3121, 0.7751, 0.5351, 0.9866]],     3
--------------------------------------------------
        [[0.6985, 0.7313, 0.4613, 0.1862],
         [0.5678, 0.4485, 0.4514, 0.7747],
         [0.2221, 0.6104, 0.2327, 0.9274],
         [0.2359, 0.2159, 0.3979, 0.2519]]])    4
---------------------------------------------------
i = a[:,:,0] > 0.5
i
tensor([[0, 1, 0, 0],
        [0, 1, 0, 0],
        [0, 1, 1, 0],
        [1, 1, 0, 0]])

b = a[i]
b
tensor([[0.7185, 0.2953, 0.6841, 0.1045],
        [0.6817, 0.4053, 0.2318, 0.1309],
        [0.8265, 0.8029, 0.3165, 0.2020],
        [0.6041, 0.0118, 0.8386, 0.5076],
        [0.6985, 0.7313, 0.4613, 0.1862],
        [0.5678, 0.4485, 0.4514, 0.7747]])

=====================================================
What I want is as below:

tensor([[[0.7185, 0.2953, 0.6841, 0.1045]],        1
------------------------------------------------------
        [[0.6817, 0.4053, 0.2318, 0.1309]],        2
-------------------------------------------------------
        [[0.8265, 0.8029, 0.3165, 0.2020],
        [0.6041, 0.0118, 0.8386, 0.5076]],         3
-----------------------------------------------------
        [[0.6985, 0.7313, 0.4613, 0.1862],
        [0.5678, 0.4485, 0.4514, 0.7747]]])        4
----------------------------------------------------

Answer 1

在第一个维度中只需 unsqueeze（或您想要的任何其他维度）。

import torch

a = torch.rand(4, 4, 4)

i = a[:, :, 0] > 0.5
b = a[i]
print(b.unsqueeze(dim=1))

给你一些类似的东西：

tensor([[[0.9476, 0.3862, 0.4544, 0.5905]],

        [[0.9413, 0.9987, 0.6411, 0.6876]],

        [[0.5807, 0.6687, 0.0952, 0.1582]],

        [[0.6057, 0.6513, 0.4329, 0.2501]],

        [[0.8998, 0.4524, 0.9219, 0.0447]]])

顺便说一句。通常不需要它，因为如果需要（并且可能），这个 2D 形状会通过广播扩展为 3D。