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)
介绍
看这一章前，请先看上一章，基于转置卷积解码的卷积自编码器的图像去噪 。
上一章建立的自编码降噪器，在解码阶段使用转置卷积放大特征映射，不是唯一的上采样方法，比如说 UpSamplingBilinear2d() 函数对特征映射进行线性插值，得到上采样的特征映射，然后用卷积层对插值后的特征映射进行修正，也可以组成自编码网络解码器。
接下来的代码建立一个新的卷积自编码网络。在网络建立部分，编码器与之前相同，但是在解码阶段做了一定调整。数据预处理，网络训练，模型验证等部分的代码和上一节相同。
扫描二维码关注公众号，回复：
13713249 查看本文章
# 传输代码文件到服务器端
scp "E:\\Jupyter WorkSpace\\PytorchLearn\\torch_autoencoder_conv(upsample).ipynb" [email protected]:/home/mist/
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 文件中读取数据的函数，并对数据进行增强。
# 定义一个将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.UpsamplingBilinear2d(scale_factor=2), # [, 256, 48, 48]
nn.Conv2d(256, 128, 3, stride=1, padding=1), # [, 128, 48, 48]
nn.ReLU(True),
nn.BatchNorm2d(128),
nn.Conv2d(128, 64, 3, stride=1, padding=1), # [, 128, 48, 48]
nn.ReLU(True),
nn.BatchNorm2d(64),
nn.UpsamplingBilinear2d(scale_factor=2), # [, 64, 96, 96]
nn.Conv2d(64, 32, 3, stride=1, padding=1), # [, 32, 96, 96]
nn.ReLU(True),
nn.BatchNorm2d(32),
nn.Conv2d(32, 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
UpsamplingBilinear2d-21 [-1, 256, 48, 48] 0
Conv2d-22 [-1, 128, 48, 48] 295,040
ReLU-23 [-1, 128, 48, 48] 0
BatchNorm2d-24 [-1, 128, 48, 48] 256
Conv2d-25 [-1, 64, 48, 48] 73,792
ReLU-26 [-1, 64, 48, 48] 0
BatchNorm2d-27 [-1, 64, 48, 48] 128
UpsamplingBilinear2d-28 [-1, 64, 96, 96] 0
Conv2d-29 [-1, 32, 96, 96] 18,464
ReLU-30 [-1, 32, 96, 96] 0
BatchNorm2d-31 [-1, 32, 96, 96] 64
Conv2d-32 [-1, 3, 96, 96] 867
Sigmoid-33 [-1, 3, 96, 96] 0
================================================================
Total params: 982,275
Trainable params: 982,275
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.11
Forward/backward pass size (MB): 89.30
Params size (MB): 3.75
Estimated Total Size (MB): 93.15
----------------------------------------------------------------
/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.494127919506788 dB
去噪后的PSNR: 23.278049519478348 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的平均提升量为: 4.911651127529821 dB
可知，去噪的效果显著。
对比上一节的，基于转置卷积解码的卷积自编码器后发现，基于线性插值的卷积自编码器上升量只有 4.9 d B 4.9dB 4.9dB ，并不如基于转置卷积的卷积自编码器上升了 5.81 d B 5.81dB 5.81dB 。
至此，我们学习了三个方向的内容，基于线性层的自编码网络对于图像进行降维，它本质上是一个自监督学习过程。基于卷积和转置卷积的卷积自编码网络，基于卷积和线性插值的卷积自编码网络，对图像进行去噪。