Рассматривать кортеж / список тензоров как единый тензор

Question

Я использую Pytorch для некоторых задач по обучению робототехнике. Я хотел бы использовать как изображения, так и информацию о состоянии в качестве наблюдений для этой задачи. Реализация, которую я использую, не поддерживает это напрямую, поэтому я делаю некоторые изменения. Ожидаемые наблюдения - это либо состояние, как одномерный тензор, либо изображения как трехмерный тензор (каналы, ширина, высота). В моей задаче я хотел бы, чтобы наблюдение было кортежем Тензоров.

Во многих местах в моей кодовой базе, конечно, ожидается, что наблюдение будет одним Tensor, а не кортежем Tensor. Есть ли простой способ рассматривать кортеж Тензор как единый Тензор?

Например, я бы хотел:

observation.to(device)

, чтобы работать как обычно, когда observation является единственным Tensor, и вызывать .to(device) для каждого Tensor, когда observation является кортежем Tensors.

Должно быть достаточно просто создать тип данных, который может это поддерживать, но мне интересно, такой тип данных уже существует? Я ничего не нашел до сих пор.

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 , чтобы пойти в другом направлении.