Pytorch: 图像自编码器-卷积自编码网络(转置卷积解码)和图像去噪
Copyright: Jingmin Wei, Pattern Recognition and Intelligent System, School of Artificial and Intelligence, Huazhong University of Science and Technology
本教程不商用，仅供学习和参考交流使用，如需转载，请联系本人。
Reference
Stacked Auto-Encoders(Denoising)
介绍
利用卷积自编码网络去噪，利用其进行图像的编码和解码，是因为卷积操作在提取图像的信息上有较好的效果，而且可以对图像中隐藏的空间信息等内容进行较好的提取。该网络可以用于去噪，分割等。
在网络的输入图像带有噪声，而输出图像则为原始的去噪图像。在编码器阶段，会经过多个卷积、池化、激活和 BN 层操作，逐渐降低每个特征映射尺寸，如此降低至 24 × 24 24\times24 24×24 ，即图像缩小为原来的 1 16 \frac{1}{16} 161 。而解码器阶段，则通过多个转置卷积，激活和 BN 层操作，逐渐将其解码为原始图像大小并且包含 3 3 3 通道，即 3 × 96 × 96 3\times96\times96 3×96×96 的图像。
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from skimage.util import random_noise
from skimage.measure import compare_psnr
from skimage.metrics import peak_signal_noise_ratio
import hiddenlayer as hl
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as Data
import torch.optim as optim
from torchvision import transforms
from torchvision.datasets import STL10
# 模型加载选择GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = torch.device('cpu')
print(device)
print(torch.cuda.device_count())
print(torch.cuda.get_device_name(0))
cuda
1
NVIDIA GeForce RTX 3070
数据预处理
使用到的图像数据集为 - STL10 ，包含三种类型的数据，分别是带有标签的训练集和验证集，分别包含 5000 5000 5000 和 8000 8000 8000 张图像，共 10 10 10 类数据。还有一个类型包含 10 10 10 万张无标签图像，均为 96 × 96 96\times96 96×96 的 RGB 图像，可以用于无监督学习。
虽然可以直接用 torchvision.datasets.STL10() 下载，但是数据大小仅 2.5GB ，且为二进制。所以最好直接到本身网址：https://cs.stanford.edu/~acoates/stl10/ 下载，并保存到指定文件夹后解压。
为了节省时间并提高训练速度，在搭建 的网络中只使用 5000 5000 5000 张图片，其中 4000 4000 4000 张作为训练集， 1000 1000 1000 张作为验证集。
首先对数据预处理，定义一个从 .bin 文件中读取数据的函数，并对数据进行增强。
我们首先把文件和数据上传至服务器(我用的是 MistGPU，有关使用方式请自行在官网查询)。注意，带空格的文件夹如果无法识别，可以加上双引号。
# 传输代码文件到服务器端
scp "E:\\Jupyter WorkSpace\\PytorchLearn\\torch_autoencoder_conv(transpose).ipynb" [email protected]:/home/mist/
# 传输数据到服务器端
scp "E:\\Jupyter WorkSpace\\PytorchLearn\\data\\stl10_binary.tar.gz" [email protected]:/home/mist/data/
# z解压文件
cd ~/data
tar xf stl10_binary.tar.gz # 解压到当前目录下
然后打开 MistGPU ，继续运行代码。
# 定义一个将bin文件处理为图像数据的函数
def read_image(data_path):
with open(data_path, 'rb') as f:
datal = np.fromfile(f, dtype=np.uint8)
# 图像[num, channels, width, height]
images = np.reshape(datal, (-1, 3, 96, 96))
# 转为RGB
images = np.transpose(images, (0, 3, 2, 1))
# 图像标准化到0-1
return images / 255.0
上述函数处理后，第一位表示图像的数量，后面表示图像的 RGB 像素值(方便 matplotlib 可视化)。
用该函数读取 STL10 数据的训练数据集 train_X.bin 程序如下：
# 读取训练集，5000张96*96*3的图像
data_path = './data/stl10_binary/train_X.bin'
images = read_image(data_path)
print('images.shape:', images.shape)
images.shape: (5000, 96, 96, 3)
接下来定义一张为图像数据添加高斯噪声的函数，为每一张图像都添加随机噪声。
# 为数据添加高斯噪声
def gaussian_noise(images, sigma):
# sigma: 噪声标准差
sigma2 = sigma**2 / (255**2) # 噪声方差
images_noisy = np.zeros_like(images) # 0矩阵初始化
for i in range(images.shape[0]):
image = images[i]
# 添加噪声
image_noise = random_noise(image, mode='gaussian', var=sigma2, clip=True)
images_noisy[i] = image_noise
return images_noisy
images_noise = gaussian_noise(images, 30)
print('images_noise:', images_noise.min(), '~', images_noise.max())
images_noise: 0.0 ~ 1.0
以上代码通过 random_noise 为每张图像添加指定方差为 sigma2 的图像噪声，并且将像素值范围处理到 0 − 1 0-1 0−1 之间。从输出可知，所有像素值最大为 1 1 1 ，最小为 0 0 0 。
下面可视化其中部分图像，对比添加噪声前后的内容
# 可视化其中部分原图像和高斯噪声图
plt.figure(figsize=(14, 8))
for i in np.arange(36):
plt.subplot(6, 12, i+1)
plt.imshow(images[i, ...])
plt.axis('off')
plt.subplot(6, 12, i+37)
plt.imshow(images_noise[i, ...])
plt.axis('off')
plt.show()
数据集构建
接下来对图像切分为训练集和验证集，并进行图像增强和数据类型标准化，处理为张量格式：
# 转为[num, channels, height, width]
data_Y = np.transpose(images, (0, 3, 2, 1)) # 原图作为Labels
data_X = np.transpose(images_noise, (0, 3, 2, 1)) # 噪声图作为Inputs
# 分割为训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(data_X, data_Y, test_size=0.2, random_state=123)
# 转为张量数据
X_train = torch.tensor(X_train, dtype=torch.float32)
X_val = torch.tensor(X_val, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32)
y_val = torch.tensor(y_val, dtype=torch.float32)
# 将X, Y转化为数据集
train_data = Data.TensorDataset(X_train, y_train)
val_data = Data.TensorDataset(X_val, y_val)
print('X_train.shape:', X_train.shape)
print('X_val.shape:', X_val.shape)
print('y_train.shape:', y_train.shape)
print('y_val.shape:', y_val.shape)
X_train.shape: torch.Size([4000, 3, 96, 96])
X_val.shape: torch.Size([1000, 3, 96, 96])
y_train.shape: torch.Size([4000, 3, 96, 96])
y_val.shape: torch.Size([1000, 3, 96, 96])
接下来使用 Data.Dataloader() 方法将二者处理为数据加载器，每个 batch 包含 32 32 32 张图像。
# 定义训练集和验证集的数据加载器
train_loader = Data.DataLoader(
dataset=train_data, # 使用的数据集
batch_size=32, # 批处理样本大小
shuffle=True, # 每次训练迭代时都打乱数据
num_workers=0 # windows上只能0个进程，linux可设置为4
)
val_loader = Data.DataLoader(
dataset=val_data,
batch_size=32,
shuffle=True,
num_workers=0
)
基于转置卷积编码的网络构建
class Denoise_AutoEncoders(nn.Module):
def __init__(self):
super(Denoise_AutoEncoders, self).__init__()
# 定义Encoder编码层
self.Encoder = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1), # [, 64, 96, 96]
nn.ReLU(True),
nn.BatchNorm2d(64),
nn.Conv2d(64, 64, 3, stride=1, padding=1), # [, 64, 96, 96]
nn.ReLU(True),
nn.BatchNorm2d(64),
nn.Conv2d(64, 64, 3, stride=1, padding=1), # [, 64, 96, 96]
nn.ReLU(True),
nn.MaxPool2d(kernel_size=2, stride=2), # [, 64, 48, 48]
nn.BatchNorm2d(64),
nn.Conv2d(64, 128, 3, stride=1, padding=1), # [, 128, 48, 48]
nn.ReLU(True),
nn.BatchNorm2d(128),
nn.Conv2d(128, 128, 3, stride=1, padding=1), # [, 128, 48, 48]
nn.ReLU(True),
nn.BatchNorm2d(128),
nn.Conv2d(128, 256, 3, stride=1, padding=1), # [, 256, 48, 48]
nn.ReLU(True),
nn.MaxPool2d(2, 2), # [, 256, 24, 24]
nn.BatchNorm2d(256),
)
# 定义Decoder解码层，使用转置卷积
self.Decoder = nn.Sequential(
nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=3, stride=1, padding=1), # [, 256, 24, 24]
nn.ReLU(True),
nn.BatchNorm2d(128),
nn.ConvTranspose2d(128, 128, 3, stride=2, padding=1, output_padding=1), # [, 128, 48, 48]
nn.ReLU(True),
nn.BatchNorm2d(128),
nn.ConvTranspose2d(128, 64, 3, stride=1, padding=1), # [, 64, 48, 48]
nn.ReLU(True),
nn.BatchNorm2d(64),
nn.ConvTranspose2d(64, 32, 3, stride=1, padding=1), # [, 32, 48, 48]
nn.ReLU(True),
nn.BatchNorm2d(32),
nn.ConvTranspose2d(32, 32, 3, stride=1, padding=1), # [, 32, 48, 48]
nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1), # [, 16, 96, 96]
nn.ReLU(True),
nn.BatchNorm2d(16),
nn.ConvTranspose2d(16, 3, 3, stride=1, padding=1), # [, 3, 96, 96]
nn.Sigmoid(),
)
def forward(self, x):
encoder = self.Encoder(x)
decoder = self.Decoder(encoder)
return encoder, decoder
上述网络定义了一个卷积自编码器。在编码层，卷积核大小均为 3 × 3 3\times3 3×3 ，激活函数为 ReLU，采用最大值池化，并使用了 Batch Normalization 。编码后，尺寸从 96 × 96 96\times96 96×96 缩小为 24 × 24 24\times24 24×24 ，并且通道数从 3 3 3 增加到 256 256 256 。
在解码层，操作相反，对特征映射做转置卷积，从而放大了特征映射，从 24 × 24 24\times24 24×24 放大到 96 × 96 96\times96 96×96 ，并且通道从 256 256 256 逐渐过渡为 3 3 3 ，对应着原始 RGB 的通道数。
接下来定义一个网络实体：
# 定义网络对象
DAEmodel = Denoise_AutoEncoders().to(device)
from torchsummary import summary
summary(DAEmodel, input_size=(3, 96, 96))
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 96, 96] 1,792
ReLU-2 [-1, 64, 96, 96] 0
BatchNorm2d-3 [-1, 64, 96, 96] 128
Conv2d-4 [-1, 64, 96, 96] 36,928
ReLU-5 [-1, 64, 96, 96] 0
BatchNorm2d-6 [-1, 64, 96, 96] 128
Conv2d-7 [-1, 64, 96, 96] 36,928
ReLU-8 [-1, 64, 96, 96] 0
MaxPool2d-9 [-1, 64, 48, 48] 0
BatchNorm2d-10 [-1, 64, 48, 48] 128
Conv2d-11 [-1, 128, 48, 48] 73,856
ReLU-12 [-1, 128, 48, 48] 0
BatchNorm2d-13 [-1, 128, 48, 48] 256
Conv2d-14 [-1, 128, 48, 48] 147,584
ReLU-15 [-1, 128, 48, 48] 0
BatchNorm2d-16 [-1, 128, 48, 48] 256
Conv2d-17 [-1, 256, 48, 48] 295,168
ReLU-18 [-1, 256, 48, 48] 0
MaxPool2d-19 [-1, 256, 24, 24] 0
BatchNorm2d-20 [-1, 256, 24, 24] 512
ConvTranspose2d-21 [-1, 128, 24, 24] 295,040
ReLU-22 [-1, 128, 24, 24] 0
BatchNorm2d-23 [-1, 128, 24, 24] 256
ConvTranspose2d-24 [-1, 128, 48, 48] 147,584
ReLU-25 [-1, 128, 48, 48] 0
BatchNorm2d-26 [-1, 128, 48, 48] 256
ConvTranspose2d-27 [-1, 64, 48, 48] 73,792
ReLU-28 [-1, 64, 48, 48] 0
BatchNorm2d-29 [-1, 64, 48, 48] 128
ConvTranspose2d-30 [-1, 32, 48, 48] 18,464
ReLU-31 [-1, 32, 48, 48] 0
BatchNorm2d-32 [-1, 32, 48, 48] 64
ConvTranspose2d-33 [-1, 32, 48, 48] 9,248
ConvTranspose2d-34 [-1, 16, 96, 96] 4,624
ReLU-35 [-1, 16, 96, 96] 0
BatchNorm2d-36 [-1, 16, 96, 96] 32
ConvTranspose2d-37 [-1, 3, 96, 96] 435
Sigmoid-38 [-1, 3, 96, 96] 0
================================================================
Total params: 1,143,587
Trainable params: 1,143,587
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.11
Forward/backward pass size (MB): 80.86
Params size (MB): 4.36
Estimated Total Size (MB): 85.33
----------------------------------------------------------------
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at /home/mist/pytorch/c10/core/TensorImpl.h:1156.)
return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
# 输出网络结构
from torchviz import make_dot
x = torch.randn(1, 3, 96, 96).requires_grad_(True)
y = DAEmodel(x.to(device))
myDAENet_vis = make_dot(y, params=dict(list(DAEmodel.named_parameters()) + [('x', x)]))
myDAENet_vis
网络训练和预测
定义优化算法为 Adam，损失函数采用均方误差，并可视化损失大小的变化。
LR = 0.0003
optimizer = optim.Adam(DAEmodel.parameters(), lr=LR) # Adam优化器
loss_fuc = nn.MSELoss().to(device) # 损失函数
# 使用Hiddenlayer可视化
historyl = hl.History()
# 使用Canvas
canvasl = hl.Canvas()
train_num = 0
val_num = 0
# 对模型迭代训练，对所有数据训练epoch轮
for epoch in range(10):
train_loss_epoch = 0
val_loss_epoch = 0
# 对训练加载器的数据迭代优化
DAEmodel.train()
for step, (b_x, b_y) in enumerate(train_loader):
b_x, b_y = b_x.to(device), b_y.to(device)
_, output = DAEmodel(b_x)
loss = loss_fuc(output, b_y) # 均方根误差
optimizer.zero_grad() # 清空过往梯度
loss.backward() # 梯度反向传播
optimizer.step() # 根据梯度更新参数
train_loss_epoch += loss.item() * b_x.size(0)
train_num = train_num + b_x.size(0)
# 对验证加载器的数据进行模型验证
DAEmodel.eval()
for step, (b_x, b_y) in enumerate(val_loader):
b_x, b_y = b_x.to(device), b_y.to(device)
_, output = DAEmodel(b_x)
loss = loss_fuc(output, b_y)
val_loss_epoch += loss.item() * b_x.size(0)
val_num = val_num + b_x.size(0)
# 计算一个epoch的损失
train_loss = train_loss_epoch / train_num
val_loss = val_loss_epoch / val_num
# 保存每个epoch的输出loss
historyl.log(epoch, train_loss=train_loss, val_loss=val_loss)
# 可视化网络训练过程
with canvasl:
canvasl.draw_plot([historyl['train_loss'], historyl['val_loss']])
评价去噪效果
下面针对验证集中的一张图像使用训练好的降噪器去噪，并与原始图像对比效果
# 输入
imageindex = 1 # 图像索引
im = X_val[imageindex, ...]
im = im.unsqueeze(0)
im_nose = np.transpose(im.data.numpy(), (0, 3, 2, 1))
im_nose = im_nose[0, ...]
# 网络去噪结果
DAEmodel.eval()
_, output = DAEmodel(im.cuda())
im_de = np.transpose(output.data.cpu().numpy(), (0, 3, 2, 1))
im_de = im_de[0, ...]
# 标签(无噪声图像)
im = y_val[imageindex, ...]
im_or = im.unsqueeze(0)
im_or = np.transpose(im_or.data.numpy(), (0, 3, 2, 1))
im_or = im_or[0, ...]
# 计算去噪后的PSNR
print('加噪后的PSNR:', peak_signal_noise_ratio(im_or, im_nose), 'dB')
print('去噪后的PSNR:', peak_signal_noise_ratio(im_or, im_de), 'dB')
加噪后的PSNR: 19.41260585144029 dB
去噪后的PSNR: 24.97685547084317 dB
上述代码是对降噪前后的图像分别计算出 PSNR，即峰值信噪比，值越大说明两个图像之间越相似，可用于表示去噪效果。
接下来对原始图像、带噪图像和去噪图像分别可视化：
# 结果可视化
plt.figure(figsize=(12, 4))
plt.subplot(1, 3, 1)
plt.imshow(im_or)
plt.axis('off')
plt.title('Original Image')
plt.subplot(1, 3, 2)
plt.imshow(im_nose)
plt.axis('off')
plt.title('Noise Image $\sigma$=30')
plt.subplot(1, 3, 3)
plt.imshow(im_de)
plt.axis('off')
plt.title('Denoising Image')
plt.show()
从图中可以看出，去噪效果显著，已经看不到图像的噪声点了，而且去噪图像非常平滑，和原始图像很相近。
下面针对整个验证数据集使用降噪器，计算出所有图像降噪前后的 PSNR 提升大小的均值，来衡量多张数据上的降噪情况。
# 计算整个验证集去噪后的PSNR提升的均值
PSNR_val = []
DAEmodel.eval()
for i in range(X_val.shape[0]):
imageindex = i
# 输入
im = X_val[imageindex, ...]
im = im.unsqueeze(0)
im_nose = np.transpose(im.data.numpy(), (0, 3, 2, 1))
im_nose = im_nose[0, ...]
# 去噪
_, output = DAEmodel(im.cuda())
im_de = np.transpose(output.data.cpu().numpy(), (0, 3, 2, 1))
im_de = im_de[0, ...]
# 输出
im = y_val[imageindex, ...]
im_or = im.unsqueeze(0)
im_or = np.transpose(im_or.data.numpy(), (0, 3, 2, 1))
im_or = im_or[0, ...]
# 计算去噪后的PSNR
PSNR_val.append(peak_signal_noise_ratio(im_or, im_de) - peak_signal_noise_ratio(im_or, im_nose))
print('PSNR的平均提升量为:', np.mean(PSNR_val), 'dB')
PSNR的平均提升量为: 5.813400503195555 dB
可知，峰值信噪比提升了 5.81 d B 5.81dB 5.81dB ，去噪的效果非常显著。
受于训练时间和设备的约束，并没有使用更多的图像以及不同类型的噪声训练，所以得到的降噪器还有一定的局限性，但是从降噪后的图像已经反映出了基于卷积神经网络的降噪自编码器在图像去噪方面的有效性。