为什么我们在图像处理时使用 Unsqueeze() 函数?

Why we use Unsqueeze() function while image processing?

我正在尝试做一个指导项目,它与图像处理有关。在进行图像处理时,讲师使用 Unsqueeze(0) 函数来设置床的尺寸。我想知道改变床尺寸后会发生什么。下面给出代码供大家参考。

非常感谢您的快速回复。

from torchvision import transforms as T

def preprocess(img_path,max_size = 500):
  image = Image.open(img_path).convert('RGB')

  if max(image.size) > max_size:
    size = max_size
  else:
    size = max(image.size)

  img_transform = T.Compose([
                             T.Resize(size),
                             T.ToTensor(),
                             T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
  ])

  image = img_transform(image)
  image = image.unsqueeze(0)
  return image

这一行之后:

image = Image.open(img_path).convert('RGB')

image 可能是某种 3D 矩阵。信息的一种排列方式是维度 [Channel, Row, Intensity],因此您有:

  • 包含许多行的 R 矩阵,每行包含红色通道的强度值
  • 包含许多行的 G 矩阵,每行包含绿色通道的强度值
  • 包含许多行的 B 矩阵,每行包含蓝色通道的强度值

现在,在机器学习中,当我们训练模型时,我们很少对只有一个示例感兴趣。我们训练批量示例。批处理只是一组堆叠在一起的图像,因此我们需要从:[Channel, Row, Intensity][Batch, Channel, Row, Intensity]

这就是 unsqueeze(0) 所做的,它添加了一个新的第零维,用于使图像可堆叠。

此处使用 unsqueeze 可能是因为您正在使用卷积神经网络。

加载图像时,它通常有 3 个维度,宽度、高度颜色通道数。对于黑白图像,颜色通道数为1,对于彩色图像,有3个颜色通道(红、绿、蓝,RGB)。因此,在您的情况下,当您加载图像并将其存储为张量时,它的形状为:

image = img_transform(image) # the resulting image has shape [3, H, W]

请注意,维度顺序为 [channels, height, width] 而不是其他顺序的原因是 PyTorch。其他 libraries/software 可能会有所不同。

但是,3 个维度对于 2D 卷积神经网络来说是不够的。在深度学习中,数据是分批处理的。因此,在卷积神经网络的情况下,它不会一次只处理一张图像,而是同时并行处理 N 张图像。我们称这个图像集合为一批。因此,您将拥有 [N, C, H, W] 而不是尺寸 [C, H, W](如 here 所示)。例如,一批 64 张彩色 100 x 100 图像,您的形状为:

[64, 3, 100, 100]

现在,如果你想一次只处理一张图片,你仍然需要把它变成批处理的形式让模型接受它。例如,如果您有一张形状为 [3, 100, 100] 的图像,您需要将其转换为 [1, 3, 100, 100]。这就是 unsqueeze(0) 所做的:

image = img_transform(image) # [3, H, W]
image = image.unsqueeze(0) # [1, 3, H, W]