基于卷积神经网络的数字重建

本文提供了一种新的方法来处理图像重建任务。

背景

最近的任务是训练一个可以实现"图-图"的神经网络。翻阅了很多论文，最后确定使用基于卷积神经网络的U-net. 因为在查找资料中，发现中文资料不多，所以将自己的经验分享出来，希望能帮到有需要的人。

方法

平台

MATLAB 2018b
使用Matlab是因为一直在使用Matlab，平常的很多作业任务都是用Matlab来实现的。而且Matlab的Deep Learning Toolbox也比较好入门。很多网络和函数都已经集成，比如今天要用到的Unet； 再者官网document非常详细，基本上可以解决你的所有问题，不过最新的很多的内容没有中文。 缺点是老版本会缺失一些新的函数和功能，这在深度学习飞速发展的今天还是有些影响的。当然，也可以自定义Layer，Loss function等等。

数据集

以MNIST的数字灰度图作为数据源A，经过一个随机生成的奇异矩阵，得到了难以识别的图像矩阵B。将数字图A作为真实图，B作为神经网络的输入，共10000对图片。
代码如下：

%%load data
%这里取巧，直接使用官方给出的测试集10000张作为训练集
images = loadMNISTImages('t10k-images.idx3-ubyte');

% 定义两个矩阵，储存数据；神经网络的输入为32x32x1的灰度图，所以需要reshape
images_N = zeros(28,28,10000);
images_32 = zeros(32,32,1,10000);

for i = 1:10000
digit = reshape(images(:, i), [28,28]);
images_N(:,:,i) = digit;
images_32(:,:,i)= imresize(digit,[32 32]);
end

处理图片，并保存

 %产生一个随机矩阵，保存，以后再生成数据还是用同一个
[M_T,~,~]= svd(rand(32,32));
% 4维数组，可以直接作为训练集
images_32P = zeros(32,32,1,10000);
for i = 1:10000
digit = images_32(:,:,:,i);
images_32P(:,:,:,i) =M_T*digit;
end 

可以用下面的代码测试

for i = 1:50
    subplot(5,10,i)
    test = images_32(:,:,:,i);
    imshow(test);
end
figure;
for i = 1:50
    subplot(5,10,i)
    test = images_32P(:,:,:,i);
    imshow(test);
end

结果为：

神经网络

使用类似于Unet的卷积神经网络。
Unet近些年在图像分割领域取得了出色的结果，也是被研究人员广泛使用。 Unet主要的目的是像素分类，最后一层是Classification Layer，并不适用于图像重建，但是这种拓扑结构在处理图像任务上非常有效。原文链接：U-Net: Convolutional Networks for Biomedical Image Segmentation
为了达到图像重建的目的，我替换了输出层，变为regression layer；并且删掉了反卷积（又叫转置卷积，上采样）时的连接层，连接层的主要任务是补充上采样时丢失的信息。 但是在此类图像重建任务中，并不需要，所以也删掉了。
最后只保留了encoder-decoder. 图中显示激活函数为ReLU，实际上是LeakyReLU

训练

在Matlab上，可以用很少的代码完成训练任务。
首先是上载数据。这里我定义的O：原图，P：经过处理后的图像。

load('digits_32_10000_P.mat');
load('digits_32_10000_O.mat');
PTrain = images_32P;
OTrain = images_32;

生成神经网络

lgraph = createUnet_regression_withoutConnection_LeakyReLU();

训练参数

%% Set the training parameters
maxEpochs = 15;
initLearningRate = 1e-4;
options = trainingOptions('adam',...
        'InitialLearnRate',initLearningRate,...
        'Shuffle','never',...
        'MaxEpochs',maxEpochs,... 
        'VerboseFrequency',5,...
        'MiniBatchSize',256,...
        'ExecutionEnvironment','parallel',...
        'Plots','training-progress');  

开始训练

%% training
modelDateTime = datestr(now,'dd-mmm-yyyy-HH-MM-SS');
net = trainNetwork(PTrain,OTrain, lgraph,options);
save(['trainedUNet-' modelDateTime '-Epoch-' num2str(maxEpochs)],'net','options');

结果

选取前20张图片，可视化。这是输入：

这是预测图：

这是原图：

由此可见，使用Unet结构的卷积神经网络在图像重建方面表现不错。
（时间关系，以后会慢慢补充完整。 有问题，请直接留言。）