pytorch - 将张量元组/列表视为单个张量

Question

我正在将 Pytorch 用于一些机器人强化学习任务。我想使用有关状态的图像和信息作为此任务的观察结果。我正在使用的实现不直接支持这一点，所以我正在做一些修改。预期的观察结果要么是状态，作为 1 维张量，要么是图像作为 3 维张量( channel ，宽度，高度)。在我的任务中，我希望观察结果是一个张量元组。

在我的代码库中的许多地方，观察当然应该是单个张量，而不是张量元组。 有没有一种简单的方法可以将张量元组视为单个张量？

例如，我想:

observation.to(device)

正常工作时 observation是单个张量，调用 .to(device)在每个张量上 observation是一个张量元组。

创建可以支持这一点的数据类型应该很简单，但我想知道这样的数据类型是否已经存在？到目前为止我还没有发现任何东西。

Answer 1

如果您的张量大小相同，则可以使用 torch.stack将它们连接成一个多维的张量。

例子:

>>> import torch
>>> a=torch.randn(2,1)
>>> b=torch.randn(2,1)
>>> c=torch.randn(2,1)
>>> a
tensor([[ 0.7691],
        [-0.0297]])
>>> b
tensor([[ 0.4844],
        [-0.9142]])
>>> c
tensor([[ 0.0210],
        [-1.1543]])
>>> torch.stack((a,b,c))
tensor([[[ 0.7691],
         [-0.0297]],

        [[ 0.4844],
         [-0.9142]],

        [[ 0.0210],
         [-1.1543]]])

然后您可以使用 torch.unbind去另一个方向。