Pytorch：与反向池化和复制填充类似的过程？

Posted 2023-03-16

【中文标题】Pytorch：与反向池化和复制填充类似的过程？

【英文标题】：Pytorch: a similar process to reverse pooling and replicate padding?

【发布时间】：2021-07-31 22:55:54

【问题描述】：

我有一个张量A，其形状为(batch_size, width, height)。假设它有这些值：

A = torch.tensor([[[0, 1],
                   [1, 0]]])

我还得到一个数字K，它是一个正整数。在这种情况下让K=2。我想做一个类似于反向池和复制填充的过程。这是预期的输出：

B = torch.tensor([[[0, 0, 1, 1],
                   [0, 0, 1, 1],
                   [1, 1, 0, 0],
                   [1, 1, 0, 0]]])

解释：对于A中的每个元素，我们将其展开为形状为(K, K)的矩阵，并放入结果张量中。我们继续对其他元素这样做，并让它们之间的步幅等于内核大小（即K）。

如何在 PyTorch 中做到这一点？目前，A 是一个二进制掩码，但如果我可以将其扩展为非二进制大小写会更好。

【参考方案1】：

平方展开

展开两次就可以得到你想要的输出：

def dilate(t, k):
  x = t.squeeze()
  x = x.unsqueeze(-1).expand([*x.shape,k])
  x = x.unsqueeze(-1).expand([*x.shape,k])
  x = torch.cat([*x], dim=1)
  x = torch.cat([*x], dim=1)
  x = x.unsqueeze(0)
  return x

B = dilate(A, k)

调整大小/最近插值

如果您不介意在较大的扩展中可能“出血”的角点（因为在确定“最近”点进行插值时，它使用欧几里得而不是曼哈顿距离），更简单的方法是 resize：

import torchvision.transforms.functional as F

B = F.resize(A, A.shape[-1]*k)

---

为了完整性：

MaxUnpool2d 将MaxPool2d 的输出（包括最大值的索引）作为输入，并计算一个部分逆运算，其中所有非最大值都设置为零。

【讨论】：

如何在通道维度复制然后应用PixelShuffle？

【参考方案2】：

你可以试试这些：

注意：以下函数将 2D 张量作为输入。如果您的张量 A 的形状为 (1, N, N)，即具有（冗余）批次/通道维度，请将 A.squeeze() 传递给 func()。

方法一：

此方法广播乘法，然后进行转置和整形操作以达到最终结果。

import torch
import torch.nn as nn

A = torch.tensor([[0, 1, 1], [1, 0, 1], [1, 1, 0]])
K = 3

def func(A, K):
    ones = torch.ones(K, K)
    tmp = ones.unsqueeze(0) * A.view(-1, 1, 1)
    tmp = tmp.reshape(A.shape[0], A.shape[1], K, K)
    res = tmp.transpose(1, 2).reshape(K * A.shape[0], K * A.shape[1])
    return res

方法二：

根据@Shai 在 cmets 中的提示，此方法在通道维度上重复 (2D) 张量K**2 次，然后使用PixelShuffle() 将行和列放大K 次。

def pixelshuffle(A, K):
    pixel_shuffle = nn.PixelShuffle(K)
    return pixel_shuffle(A.unsqueeze(0).repeat(K**2, 1, 1).unsqueeze(0)).squeeze(0).squeeze(0)

由于 nn.PixelShuffle() 仅将 4D 张量作为输入，因此需要在 repeat() 之后解压。另请注意，由于从nn.PixelShuffle() 返回的张量也是 4D，因此两个 squeeze()s 紧随其后以确保我们得到一个 2D 张量作为输出。

一些示例输出：

A = torch.tensor([[0, 1], [1, 0]])
func(A, 2)
# tensor([[0., 0., 1., 1.],
#         [0., 0., 1., 1.],
#         [1., 1., 0., 0.],
#         [1., 1., 0., 0.]])

---

pixelshuffle(A, 2)
# tensor([[0, 0, 1, 1],
#         [0, 0, 1, 1],
#         [1, 1, 0, 0],
#         [1, 1, 0, 0]])

请随时要求进一步澄清，让我知道它是否适合您。

基准测试：

我将我的答案func() 和pixel shuffle() 与上面@iacob 的dilate() 函数进行了基准测试，发现我的答案稍快。

A = torch.randint(3, 100, (20, 20))
assert (dilate(A, 5) == func(A, 5)).all()
assert (dilate(A, 5) == pixelshuffle(A, 5)).all()

%timeit dilate(A, 5)
# 142 µs ± 2.54 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit func(A, 5)
# 57.9 µs ± 1.67 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit pixelshuffle(A, 5)
# 81.6 µs ± 970 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)