利用 TensorFlow2 实现去噪自编码器：去除图像中的噪声

2023-10-20 18:16:26

在计算机视觉和深度学习领域，图像去噪是一个重要的课题，去噪自编码器 (Denoising Autoencoder, DAE) 就是一种旨在去除图像中噪声的神经网络模型。

在本文中，我们将使用 TensorFlow2 来实现一个 DAE，并利用 MNIST 数据集来评估其性能。

1. 数据准备

首先，我们需要加载并预处理 MNIST 数据集。MNIST 数据集是一个包含 70,000 张手写数字图像的公共数据集，其中包含 60,000 张训练图像和 10,000 张测试图像。

import tensorflow as tf

# 加载 MNIST 数据集
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

# 归一化图像
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0

# 添加噪声
x_train_noisy = x_train + 0.2 * tf.random.normal(shape=x_train.shape)
x_test_noisy = x_test + 0.2 * tf.random.normal(shape=x_test.shape)

# 将图像转换为张量
x_train_noisy = tf.data.Dataset.from_tensor_slices(x_train_noisy).batch(32)
x_test_noisy = tf.data.Dataset.from_tensor_slices(x_test_noisy).batch(32)

2. 模型架构

我们的 DAE 模型将包含以下层：

输入层：接受噪声图像作为输入。
编码器：将输入图像压缩成一个较小的中间表示。
解码器：将中间表示解压缩回重建图像。

class DenoisingAutoencoder(tf.keras.Model):

    def __init__(self):
        super(DenoisingAutoencoder, self).__init__()

        # 编码器
        self.encoder = tf.keras.Sequential([
            tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
            tf.keras.layers.MaxPooling2D((2, 2), padding='same'),
            tf.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
            tf.keras.layers.MaxPooling2D((2, 2), padding='same')
        ])

        # 中间表示层
        self.latent_dim = 128
        self.latent_layer = tf.keras.layers.Dense(self.latent_dim, activation='relu')

        # 解码器
        self.decoder = tf.keras.Sequential([
            tf.keras.layers.Dense(784, activation='relu'),
            tf.keras.layers.Reshape((28, 28, 1)),
            tf.keras.layers.Conv2DTranspose(32, (3, 3), activation='relu', padding='same'),
            tf.keras.layers.UpSampling2D((2, 2)),
            tf.keras.layers.Conv2DTranspose(32, (3, 3), activation='relu', padding='same'),
            tf.keras.layers.UpSampling2D((2, 2)),
            tf.keras.layers.Conv2D(1, (3, 3), activation='sigmoid', padding='same')
        ])

    def call(self, inputs):

        # 编码器
        x = self.encoder(inputs)

        # 中间表示层
        latent = self.latent_layer(x)

        # 解码器
        outputs = self.decoder(latent)

        return outputs

3. 训练模型

现在我们可以训练我们的 DAE 模型了。我们将使用二叉交叉熵损失函数和 Adam 优化器。

model = DenoisingAutoencoder()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

# 损失函数
loss_fn = tf.keras.losses.BinaryCrossentropy()

# 训练循环
epochs = 10

for epoch in range(epochs):

    for step, (x_batch_noisy, _) in enumerate(x_train_noisy):

        with tf.GradientTape() as tape:

            # 前向传播
            reconstructed_images = model(x_batch_noisy)

            # 计算损失
            loss = loss_fn(x_batch_noisy, reconstructed_images)

        # 反向传播
        gradients = tape.gradient(loss, model.trainable_variables)

        # 更新权重
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # 打印训练信息
    print(f'Epoch: {epoch + 1}, Loss: {loss.numpy():.4f}')

4. 评估模型

在训练模型后，我们可以使用测试集来评估其性能。

# 评估模型
reconstructed_images = model.predict(x_test_noisy)

# 计算重建误差
reconstruction_error = np.mean(np.square(x_test - reconstructed_images))

# 打印评估结果
print(f'Reconstruction Error: {reconstruction_error:.4f}')