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

Question

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

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

例如，我想：

observation.to(device)

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

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

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 去另一个方向。