如何在 PyTorch 中 combine/stack 张量和组合维度？

Question

我需要将 4 张量，代表灰度图像，大小 [1,84,84]，组合成一堆形状 [4,84,84]，代表四个灰度图像，每个图像表示为 "channel" 张量样式 CxWxH.

我正在使用 PyTorch。

我试过使用 torch.stack 和 torch.cat，但如果其中之一是解决方案，我就无法找到正确的 prep/methodology 来获得我的结果。

感谢您的帮助。

import torchvision.transforms as T

class ReplayBuffer:
    def __init__(self, buffersize, batchsize, framestack, device, nS):
        self.buffer = deque(maxlen=buffersize)
        self.phi = deque(maxlen=framestack)
        self.batchsize = batchsize
        self.device = device

        self._initialize_stack(nS)

    def get_stack(self):
        #t =  torch.cat(tuple(self.phi),dim=0)
        t =  torch.stack(tuple(self.phi),dim=0)
        return t

    def _initialize_stack(self, nS):
        while len(self.phi) < self.phi.maxlen:
            self.phi.append(torch.tensor([1,nS[1], nS[2]]))

a = ReplayBuffer(buffersize=50000, batchsize=64, framestack=4, device='cuda', nS=[1,84,84])
print(a.phi)
s = a.get_stack()
print(s, s.shape)

以上代码returns:

print(a.phi)

deque([tensor([ 1, 84, 84]), tensor([ 1, 84, 84]), tensor([ 1, 84, 84]), tensor([ 1, 84, 84])], maxlen=4)

print(s, s.shape)

tensor([[ 1, 84, 84],
        [ 1, 84, 84],
        [ 1, 84, 84],
        [ 1, 84, 84]]) torch.Size([4, 3])

但我想要的是 return 只是 [4, 84, 84]。我怀疑这很简单，但它正在逃避我。

Answer 1

看来你误解了torch.tensor([1, 84, 84])的意思。一起来看看：

torch.tensor([1, 84, 84])
print(x, x.shape) #tensor([ 1, 84, 84]) torch.Size([3])

从上面的例子可以看出，它给了你一个只有一维的张量。

根据你的问题陈述，你需要一个形状为 [1,84,84] 的张量。它是这样的：

from collections import deque
import torch
import torchvision.transforms as T

class ReplayBuffer:
    def __init__(self, buffersize, batchsize, framestack, device, nS):
        self.buffer = deque(maxlen=buffersize)
        self.phi = deque(maxlen=framestack)
        self.batchsize = batchsize
        self.device = device

        self._initialize_stack(nS)

    def get_stack(self):
        t =  torch.cat(tuple(self.phi),dim=0)
#         t =  torch.stack(tuple(self.phi),dim=0)
        return t

    def _initialize_stack(self, nS):
        while len(self.phi) < self.phi.maxlen:
#             self.phi.append(torch.tensor([1,nS[1], nS[2]]))
            self.phi.append(torch.zeros([1,nS[1], nS[2]]))

a = ReplayBuffer(buffersize=50000, batchsize=64, framestack=4, device='cuda', nS=[1,84,84])
print(a.phi)
s = a.get_stack()
print(s, s.shape)

请注意，torch.cat 为您提供形状为 [4, 84, 84] 的张量，而 torch.stack 为您提供形状为 [4, 1, 84, 84] 的张量。它们的区别可以在

找到

如何在 PyTorch 中 combine/stack 张量和组合维度？

How to combine/stack tensors and combine dimensions in PyTorch?

python

stack

concatenation

pytorch

tensor