如何用PyTorch为MNIST数据集添加可逆噪声？

我尝试为MNIST数据集添加可逆噪声用于实验，期望实现「带噪图像 - 噪声 = 原始图像」的可逆效果，但实际操作后无法还原。以下是我的实现代码及生成的效果图：

原实现代码 #

import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
from PIL import Image
import torchvision

def display_img(pixels, label = None):
    plt.imshow(pixels, cmap="gray")
    if label:    
        plt.title("Label: %d" % label)
    plt.axis("off")
    plt.show()

class NoisyMNIST(torchvision.datasets.MNIST):
    def __init__(self, root, train=True, transform=None, target_transform=None, download=False):
        super(NoisyMNIST, self).__init__(root, train=train, transform=transform, target_transform=target_transform, download=download)

    def __getitem__(self, index):
        img, target = self.data[index], self.targets[index]
        img = Image.fromarray(img.numpy(), mode="L")

        if self.transform is not None:
            img = self.transform(img)
        
        # add the noise
        noise_level = 0.3
        noise = self.generate_safe_random_tensor(img) * noise_level
        noisy_img = img + noise
        
        return noisy_img, noise, img, target

    def generate_safe_random_tensor(self, img):
        """generates random noise for an image but limits the pixel values between -1 and 1""" 
       
        min_values = torch.clamp(-1 - img, max=0)
        max_values = torch.clamp(1 - img, min=0)
       
        return torch.rand(img.shape) * (max_values - min_values) + min_values

# Define transformations to apply to the data
transform = transforms.Compose([
    transforms.ToTensor(),  # Convert images to tensors
    transforms.Normalize((0.1307,), (0.3081,)),
])

train_dataset = NoisyMNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = NoisyMNIST(root='./data', train=False, download=True, transform=transform)

np_noise = train_dataset[img_id][1]
np_data = train_dataset[img_id][0]

display_img(np_data_sub_noise, 4)

原生成效果图 #

• 原图：
  
• 带噪图：
  
• 减噪后图（错误效果）：
  

问题原因与修复方案 #

核心问题 #

1. 缺失反归一化处理：原代码直接展示归一化后的张量（数值范围非[0,1]），matplotlib自动缩放导致视觉上无法匹配原图，但实际张量数值是可逆的。
2. 代码环节缺失：未定义img_id、未正确计算减噪图像，也未处理张量的通道维度（ToTensor后为(1,28,28)，需转为(28,28)才能展示）。

修复后代码 #

import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from PIL import Image
import torchvision
import matplotlib.pyplot as plt
import numpy as np

def display_img(pixels, label=None):
    plt.imshow(pixels, cmap="gray")
    if label:    
        plt.title(f"Label: {label}")
    plt.axis("off")
    plt.show()

class NoisyMNIST(torchvision.datasets.MNIST):
    def __init__(self, root, train=True, transform=None, target_transform=None, download=False):
        super().__init__(root, train=train, transform=transform, target_transform=target_transform, download=download)

    def __getitem__(self, index):
        img, target = self.data[index], self.targets[index]
        img = Image.fromarray(img.numpy(), mode="L")

        if self.transform is not None:
            img = self.transform(img)
        
        # 添加噪声
        noise_level = 0.3
        noise = self.generate_safe_random_tensor(img) * noise_level
        noisy_img = img + noise
        
        return noisy_img, noise, img, target

    def generate_safe_random_tensor(self, img):
        """生成随机噪声，确保加噪后像素值在[-1, 1]之间"""
        min_values = torch.clamp(-1 - img, max=0)
        max_values = torch.clamp(1 - img, min=0)
        return torch.rand(img.shape) * (max_values - min_values) + min_values

# 数据变换
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])

train_dataset = NoisyMNIST(root='./data', train=True, download=True, transform=transform)

# 选择样本并计算减噪图像
img_id = 0
noisy_img, noise, orig_img, target = train_dataset[img_id]
denoised_img = noisy_img - noise

# 辅助函数：将归一化张量转为可展示的numpy数组
def tensor_to_display(tensor):
    mean, std = 0.1307, 0.3081
    # 反归一化+维度压缩+范围裁剪
    np_arr = tensor.squeeze().numpy() * std + mean
    np_arr = np.clip(np_arr, 0, 1)
    return np_arr

# 验证数值可逆性
print(f"张量数值是否一致：{torch.allclose(denoised_img, orig_img, atol=1e-6)}")

# 展示图像
display_img(tensor_to_display(orig_img), target)
display_img(tensor_to_display(noisy_img), target)
display_img(tensor_to_display(denoised_img), target)

修复效果 #

修复后：

• 数值层面：denoised_img与orig_img张量完全一致（可通过torch.allclose验证）
• 视觉层面：减噪后的图像与原图完全匹配

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