在 PyTorch 中将 5D 张量转换为 4D 张量

Question

In PyTorch I have a 5D tensor X of dimensions B x 9 x C x H x W . I want to convert it into a 4D tensor Y with dimensions B x 9C x H x W such that concatenation happens channel wise.

To illustrate let,

a = X[1,0,:,:,:]
b = X[1,1,:,:,:]
c = X[1,2,:,:,:]
...
i = X[1,8,:,:,:]

Then in the tensor Y , a to i should be channel wise concatenated.

Answer 1

You can easily broadcast to a new shape with torch.reshape :

b, n, c, h, w = X.shape
X = X.reshape(b, n*c, h, w)

Answer 2

Try:

import torch
x = torch.rand(3, 4, 3, 2, 6)
print(x.shape)
y=x.flatten(start_dim=1, end_dim=2)
print(y.shape)

---

torch.Size([3, 4, 3, 2, 6])
torch.Size([3, 12, 2, 6])