PyTorch中如何将4D张量转换为按通道提取的(8,3)格式1D向量张量

你可以通过调整张量维度顺序后展平的方式实现需求，具体操作如下：

完整代码实现 #

import torch

# 原始输入张量
raw_tensor = torch.tensor([[[[ 1.,  2.],
          [ 3.,  4.]],
         [[ 5.,  6.],
          [ 7.,  8.]],
         [[ 9., 10.],
          [11., 12.]]],
        [[[13., 14.],
          [15., 16.]],
         [[17., 18.],
          [19., 20.]],
         [[21., 22.],
          [23., 24.]]]])

# 核心转换逻辑
target_tensor = raw_tensor.permute(0, 2, 3, 1).reshape(-1, 3)
# 若需要输出为整型，可补充类型转换：
# target_tensor = raw_tensor.permute(0, 2, 3, 1).reshape(-1, 3).to(torch.int)

逻辑说明 #

1. 首先调用permute(0, 2, 3, 1)调整张量维度顺序：将原始(batch_size, channels, height, width)的维度顺序调整为(batch_size, height, width, channels)，调整后每个空间坐标（height、width对应位置）的最后一维就是该位置3个通道的数值，符合按空间位置提取多通道元素的需求。
2. 再调用reshape(-1, 3)将前三个维度（batch_size、height、width）展平为一个维度，总元素数为2 * 2 * 2 = 8，最终得到形状为(8, 3)的目标张量。

输出验证 #

打印target_tensor即可得到你需要的结果：

tensor([[ 1.,  5.,  9.],
        [ 2.,  6., 10.],
        [ 3.,  7., 11.],
        [ 4.,  8., 12.],
        [13., 17., 21.],
        [14., 18., 22.],
        [15., 19., 23.],
        [16., 20., 24.]])

内容的提问来源于stack exchange，提问作者phil