PyTorch DataLoader返回图像预处理还原结果异常问题排查

我在PyTorch中创建了一个Dataset用于加载图像并输入CNN。在Dataset的__getitem__方法中，图像会先经过基于ImageNet均值的归一化预处理后返回。但发现一个问题：在Dataset内部还原预处理后的图像能得到预期的原图，但对DataLoader返回的图像执行相同还原操作时，结果却与原图不同。DataLoader不只是将预处理后的图像堆叠在一起吗？

示例代码如下（需指定本地图像路径）：

import cv2
import numpy as np
import torch


class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, img_path):
        # RGB image
        self.img = cv2.imread(img_path)[:, :, ::-1]

    def __len__(self):
        return 100

    def transform(self, img: np.array, data_type=np.float32):
        img = img.astype(data_type) / 255
        img[:, :, 0] = (img[:, :, 0] - 0.485) / 0.229
        img[:, :, 1] = (img[:, :, 1] - 0.456) / 0.224
        img[:, :, 2] = (img[:, :, 2] - 0.406) / 0.225
        return img

    def __getitem__(self, idx):
        images_raw = [self.img] * 5

        # transform images
        images_transformed = np.array([self.transform(img) for img in images_raw])
        images_transformed = torch.tensor(images_transformed)

        # test reconverting image and saving it here
        image_reconverted = transform_reverse(images_transformed[0].numpy())
        cv2.imwrite("puppy_00_re_tranformed_in_dataloader.jpg", image_reconverted[:, :, ::-1])

        return images_transformed


def main():
    img_path = "puppy.jpg"

    custom_dataset = CustomDataset(img_path)
    dataloader = torch.utils.data.DataLoader(custom_dataset, batch_size=10, shuffle=True, drop_last=True)

    for batch_data in dataloader:
        # take one of the images returned from Dataloader and apply reverse transform
        re_converted_img = transform_reverse(batch_data[0][0].numpy())
        cv2.imwrite("puppy_01_re_transformed_in_main.jpg", re_converted_img[:, :, ::-1])


def transform_reverse(img):
    img[:, :, 0] = img[:, :, 0] * 0.229 + 0.485
    img[:, :, 1] = img[:, :, 1] * 0.224 + 0.456
    img[:, :, 2] = img[:, :, 2] * 0.225 + 0.406
    img = np.round(img * 255).astype(np.uint8)
    return img


if __name__ == "__main__":
    main()

• 原图：一只浅棕色的小狗趴在白色背景上，耳朵下垂，表情温和
• DataLoader返回图像还原后的异常图：颜色严重失真，呈现偏蓝绿的怪异色调，细节模糊

问题出在**transform_reverse函数的原地修改操作，以及PyTorch张量与numpy数组的内存共享机制**：

1. 当在Dataset的__getitem__中调用images_transformed[0].numpy()时，CPU张量转numpy数组默认是共享内存的，即numpy数组和原张量指向同一块内存空间。
2. transform_reverse函数里的img[:, :, 0] = ...属于原地赋值操作，会直接修改numpy数组的内容，同时也会同步修改原张量images_transformed[0]的数值——因为二者共享内存。
3. 被修改后的张量会被返回给DataLoader，所以在main函数中对DataLoader返回的图像执行还原时，实际是对已经被“预处理+部分还原”的张量再次还原，最终得到异常结果。

解决方法 #

有两种方式可以修复：

方式1：调用还原函数前先拷贝数组

修改两处调用transform_reverse的代码：

• Dataset的__getitem__中：

  image_reconverted = transform_reverse(images_transformed[0].numpy().copy())
  

• main函数中：

  re_converted_img = transform_reverse(batch_data[0][0].numpy().copy())
  

方式2：修改还原函数，先创建数组副本再操作

直接修改transform_reverse，避免原地修改原数组：

def transform_reverse(img):
    # 先创建输入数组的副本，不修改原数据
    img = img.copy()
    img[:, :, 0] = img[:, :, 0] * 0.229 + 0.485
    img[:, :, 1] = img[:, :, 1] * 0.224 + 0.456
    img[:, :, 2] = img[:, :, 2] * 0.225 + 0.406
    img = np.round(img * 255).astype(np.uint8)
    return img

两种方式都能保证原张量的数值不被还原操作修改，DataLoader拿到的是正确的预处理后张量，还原后就能得到和Dataset内部一致的原图。

内容的提问来源于stack exchange，提问作者Carlo