在 Pytorch 中处理内存中的图像数据答案

【问题标题】：Working with image data from memory in Pytorch在 Pytorch 中处理内存中的图像数据
【发布时间】：2022-05-13 13:11:56
【问题描述】：

问题：

如果我在内存中有一个带有暗淡 (n, height, width, channels) 的数组，并且我想要一个 Pytorch 分类器将它们向前馈送并给我一个数组，其中包含数组中每个图像的类预测，我该怎么做?

背景：

我正在处理一个计算机视觉问题，我使用预先存在的代码修改了一些图像，并希望将修改后的图像发送到 Pytorch 分类器 CNN（不是我开发或控制的）。比起 Pytorch，我更习惯于 Tensorflow/Keras。

使用 Tensorflow/Keras 模型，您可以在 numpy 数组中为它们提供一堆图像，它会继续并通过模型向前馈送它们。

PS：

一位同事建议先将所有图像保存到磁盘，然后使用 DataLoader 将它们读入，但当我已经将图像保存在内存中时，这是不必要的。

对不起，如果这是一个愚蠢的问题，我试图在其他地方找到解决方案，但显然没有取得多大成功。

【问题讨论】：

你真的应该尝试提供源代码，最好是minimal reproducible example。 -- 你有 NHWC 订单。我认为 pytorch 更喜欢 NCHW。有办法转换。那么你需要制作一个 pytorch tensor （你有一个 numpy 数组吗？）。然后你把它粘在网络上（forward 打电话？）。

标签： python pytorch computer-vision dataset

【解决方案1】：

您可以创建一个自定义 DataLoader 函数，该函数将图像放入内存并返回张量，这些张量可以直接馈送到模型，而无需先将它们保存在磁盘上。一个非常简单的实现可以是：

def images_to_tensor(images):
    #images is numpy array of shape N,H,W,C
    #normalizes images between -1 and 1, comment this if you want to normalize images between 0 and 1
    images = (images.astype(np.float32) - 127.5)/128.0
    #to normalize image from 0 to 1 uncomment the line below
    #images = (images.astype(np.float32))/255.0
    #changes numpy array to tensors
    images = torch.from_numpy(images).permute(0, 3, 1, 2)
    #to convert cpu tensors to cuda uncomment the line below
    #images = images.to("cuda")
    return images

然后您可以使用该函数将图像转换为张量并将它们传递给分类模型以获得输出预测。

【讨论】：