如何在保持特定顺序的同时将展平的张量/数组转换回原始大小

Question

我有一个扁平的 pytorch 张量，它表示在获得向后传递的梯度后恢复 MaxPool2d 操作的索引。问题是，大小根据 MaxPool 前向传递之前的 kernel_size 和原始输入高度/宽度而变化。 例如，假设我有一个大小为 2x4 的输入，kernel_size 为 2x2（内核始终为正方形，步幅始终与内核相同）：

input1 = tensor([[1,2,5,6], [3,4,7,8]])

对于 input1，前向传递后，索引将为 [3,3]（从 0 开始索引）。 在向后传递中，在展平并添加梯度之后，我最终得到如下张量： 张量([0,0,0,(梯度4),0,0,0,(梯度8)])

问题是我现在需要将张量转换回原始形状，同时保持数字顺序。使用像tensor.view(input_height,input_width)这样简单的东西是行不通的，因为数字的顺序混乱了：

input1.view(2,4) = tensor([[1,2,3,4],  [5,6,7,8]])

我尝试过将其分成 input_height/input_width 组，但随后我遇到了使其适用于不同尺寸的问题。我认为有一个简单的解决方案，我只是缺乏 pytorch 或 numpy 技能来解决它:(

编辑： 好吧，我已经尝试过

input1.view(2,4)

，正如我所解释的，它的顺序不正确。并且通过排列或类似函数不可能获得正确的顺序。我尝试过这样分块：

input1.view(-1, 2)

它将张量分成 2 个块，但随后我无法将它们正确地堆叠回多个列。我到处寻找，但不知道如何进展。我什至（无济于事）问过 chatGPT 哈哈。

Answer 1

这是一种在保持顺序的情况下将展平张量重新排列回原始形状的方法：

import torch

flattened = torch.tensor([0,0,0,4,0,0,0,8]) # flattened tensor
kernel_size = 2
input_h = 2  
input_w = 4

unflatten = []
for i in range(0, len(flattened), input_w):
    unflatten.append(flattened[i:i+input_w])

result = torch.stack(unflatten)
print(result)

关键思想：

• 以大小为

  input_w

  （原始输入宽度）的块迭代展平张量
• 将这些块作为单独的张量附加以创建列表
• 将该列表堆叠回二维张量

这利用了对展平张量进行切片并将各个部分缝合在一起的优势。

示例数据的输出是：

张量([[0, 0, 0, 4], [0, 0, 0, 8]])