machine-learning - Pytorch 展开和折叠 : How do I put this image tensor back together again?

Question

我正在尝试使用展开来过滤大小为 256x256 的单 channel 2D 图像，以创建重叠为 8 的 16x16 块。如下所示:

# I = [256, 256] image
kernel_size = 16
stride = bx/2
patches = I.unfold(1, kernel_size, 
int(stride)).unfold(0, kernel_size, int(stride)) # size = [31, 31, 16, 16]

 

我已经开始尝试将图像与 fold 重新组合在一起，但我还没有完全到位。我尝试使用 View 让图像“适合”它应该的方式，但我不知道这将如何保留原始图像。也许我想多了。

# patches.shape = [31, 31, 16, 16]
patches = = filt_data_block.contiguous().view(-1, kernel_size*kernel_size) # [961, 256]
patches = patches.permute(1, 0) # size = [951, 256]

任何帮助将不胜感激。非常感谢。

Answer 1

我相信您会从使用 torch.nn.functional.fold 中受益和 torch.nn.functional.unfold在这种情况下，因为这些函数是专门为图像(或任何 4D 张量，即形状为 B X C X H X W)构建的。
让我们从展开图像开始:

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
from sklearn.datasets import load_sample_image #Used to load a sample image

dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
#Load a flower image from sklearn.datasets, crop it to shape 1 X 3 X 256 X 256:
I = torch.from_numpy(load_sample_image('flower.jpg')).permute(2,0,1).unsqueeze(0).type(dtype)[...,128:128+256,256:256+256]
kernel_size = 16
stride = kernel_size//2 
I_unf = F.unfold(I, kernel_size, stride=stride)

在这里，我们通过使用 F.unfold 获得步幅为 8 的所有 16x16 图像块。功能。这将产生形状为 torch.Size([1, 768, 961]) 的 3D 张量.即 - 961 个补丁，每个补丁有 768 = 16 X 16 X 3 个像素。
现在，假设我们希望将其折叠回 I:

I_f = F.fold(I_unf,I.shape[-2:],kernel_size,stride=stride)
norm_map = F.fold(F.unfold(torch.ones(I.shape).type(dtype),kernel_size,stride=stride),I.shape[-2:],kernel_size,stride=stride)
I_f /= norm_map

我们使用 F.fold我们告诉它 I 的原始形状, kernel_size我们曾经展开和stride用过的。折叠后I_unf我们将得到一个总和 有重叠 .这意味着生成的图像将显得饱和。因此，我们需要计算归一化映射，该归一化映射将归一化由于重叠引起的像素的多个总和。一种有效地做到这一点的方法是采用一个张量并使用 unfold其次是 fold - 模仿重叠的总和。这为我们提供了归一化映射 I_f恢复 I .
现在，我们希望绘制 I_f和 I证明内容被保留:

#Plot I:
plt.imshow(I[0,...].permute(1,2,0).cpu()/255)

#Plot I_f:
plt.imshow(I_f[0,...].permute(1,2,0).cpu()/255)

整个过程也适用于单 channel 图像。需要注意的一件事是，如果图像的空间维度不能被步幅整除，您将得到 norm_map由于某些像素无法到达而带有零(在边缘)，但您也可以轻松处理这种情况。