训练和部署用于语义分割的全卷积网络

此示例使用:

此示例说明如何使用 GPU Coder™ 在 NVIDIA® GPU 上训练和部署全卷积语义分割网络。

语义分割网络对图像中的每个像素进行分类，从而生成按类分割的图像。语义分割的应用包括自动驾驶中的道路分割以及医疗诊断中的癌细胞分割。要了解详细信息，请参阅Getting Started with Semantic Segmentation Using Deep Learning (Computer Vision Toolbox)。

为了说明训练过程，此示例将训练 FCN-8s [1]，这是一种专门用于语义图像分割的卷积神经网络 (CNN)。其他类型的用于语义分割的网络包括全卷积网络，如 SegNet 和 U-Net。您也可以将此训练过程应用于这些网络。

此示例使用剑桥大学的 CamVid 数据集 [2] 进行训练。此数据集是包含驾驶时获得的街道级视图的图像集合。该数据集提供了 32 个语义类的像素级标签，包括汽车、行人和道路。

第三方前提条件

必需

支持 CUDA® 的 NVIDIA GPU 和兼容驱动程序。

可选

NVIDIA CUDA 工具包。
NVIDIA cuDNN 库。
编译器和库的环境变量。有关支持的编译器和库的版本的信息，请参阅Third-Party Hardware (GPU Coder)。有关设置环境变量的信息，请参阅Setting Up the Prerequisite Products (GPU Coder)。

验证 GPU 环境

使用 coder.checkGpuInstall (GPU Coder) 函数验证运行此示例所需的编译器和库是否已正确设置。

envCfg = coder.gpuEnvConfig("host");
envCfg.DeepLibTarget = "cudnn";
envCfg.DeepCodegen = 1;
envCfg.Quiet = 1;
coder.checkGpuInstall(envCfg);

设置

此示例创建的全卷积语义分割网络具有从 VGG-16 网络初始化的权重。vgg16 函数检查是否存在 Deep Learning Toolbox™ Model for VGG-16 Network 支持包，并返回预训练 VGG-16 模型。

vgg16();

下载 FCN 的预训练版本。借助此预训练模型，您无需等待训练完成，即可运行整个示例。doTraining 标志控制示例是使用示例的训练网络还是预训练的 FCN 网络来生成代码。

doTraining = false;
if ~doTraining
    pretrainedURL = "https://www.mathworks.com/supportfiles" + ...
        "/gpucoder/cnn_models/fcn/FCN8sCamVid.mat";
    disp("Downloading pretrained FCN (448 MB)...");
    websave("FCN8sCamVid.mat",pretrainedURL);
end

Downloading pretrained FCN (448 MB)...

下载 CamVid 数据集

从以下 URL 下载 CamVid 数据集。

imageURL = "http://web4.cs.ucl.ac.uk/staff/g.brostow/" + ...
    "MotionSegRecData/files/701_StillsRaw_full.zip";
labelURL = "http://web4.cs.ucl.ac.uk/staff/g.brostow/" + ...
    "MotionSegRecData/data/LabeledApproved_full.zip";

outputFolder = fullfile(pwd,"CamVid");

if ~exist(outputFolder, "dir")
   
    mkdir(outputFolder)
    labelsZip = fullfile(outputFolder,"labels.zip");
    imagesZip = fullfile(outputFolder,"images.zip");   
    
    disp("Downloading 16 MB CamVid dataset labels..."); 
    websave(labelsZip, labelURL);
    unzip(labelsZip, fullfile(outputFolder,"labels"));
    
    disp("Downloading 557 MB CamVid dataset images...");  
    websave(imagesZip, imageURL);       
    unzip(imagesZip, fullfile(outputFolder,"images"));    
end

数据下载时间取决于您的 Internet 连接。下载操作完成后，示例执行才会继续。您也可以使用 Web 浏览器先将数据集下载到本地磁盘。然后，使用 outputFolder 变量指向所下载文件的位置。

加载 CamVid 图像

使用 imageDatastore 加载 CamVid 图像。通过 imageDatastore 可将大量图像高效加载至磁盘。

imgDir = fullfile(outputFolder,"images","701_StillsRaw_full");
imds = imageDatastore(imgDir);

显示其中一个图像。

I = readimage(imds,25);
I = histeq(I);
imshow(I)

加载 CamVid 像素标注图像

使用 pixelLabelDatastore (Computer Vision Toolbox) 加载 CamVid 像素标注图像数据。pixelLabelDatastore 将像素标签数据和标签 ID 封装到类名映射中。

按照 SegNet 论文 [3] 中描述的训练方法，将 CamVid 中的 32 个原始类分组为 11 个类。指定这些类。

classes = [
    "Sky"
    "Building"
    "Pole"
    "Road"
    "Pavement"
    "Tree"
    "SignSymbol"
    "Fence"
    "Car"
    "Pedestrian"
    "Bicyclist"
    ];

要将 32 个类减少为 11 个类，需要将原始数据集中的多个类组合在一起。例如，"Car" 是 "Car"、"SUVPickupTruck"、"Truck_Bus"、"Train" 和 "OtherMoving" 的组合。使用 camvidPixelLabelIDs 支持函数返回分组的标签 ID。

labelIDs = camvidPixelLabelIDs();

使用类和标签 ID 创建 pixelLabelDatastore。

labelDir = fullfile(outputFolder,"labels");
pxds = pixelLabelDatastore(labelDir,classes,labelIDs);

读取一个像素标注图像，并将其叠加在图像上方显示。

C = readimage(pxds,25);
cmap = camvidColorMap;
B = labeloverlay(I,C,"ColorMap",cmap);
imshow(B)
pixelLabelColorbar(cmap,classes);

没有颜色叠加的区域没有像素标签，在训练过程中不被使用。

分析数据集统计信息

要查看 CamVid 数据集中类标签的分布，请使用 countEachLabel (Computer Vision Toolbox)。此函数按类标签计算像素数。

tbl = countEachLabel(pxds)

tbl=11×3 table
         Name         PixelCount    ImagePixelCount
    ______________    __________    _______________

    {'Sky'       }    7.6801e+07      4.8315e+08   
    {'Building'  }    1.1737e+08      4.8315e+08   
    {'Pole'      }    4.7987e+06      4.8315e+08   
    {'Road'      }    1.4054e+08      4.8453e+08   
    {'Pavement'  }    3.3614e+07      4.7209e+08   
    {'Tree'      }    5.4259e+07       4.479e+08   
    {'SignSymbol'}    5.2242e+06      4.6863e+08   
    {'Fence'     }    6.9211e+06       2.516e+08   
    {'Car'       }    2.4437e+07      4.8315e+08   
    {'Pedestrian'}    3.4029e+06      4.4444e+08   
    {'Bicyclist' }    2.5912e+06      2.6196e+08

按类可视化像素计数。

frequency = tbl.PixelCount/sum(tbl.PixelCount);

bar(1:numel(classes),frequency)
xticks(1:numel(classes)) 
xticklabels(tbl.Name)
xtickangle(45)
ylabel("Frequency")

理想情况下，所有类都有相同数量的观测值。CamVid 中的类是不平衡的，这是街景汽车数据集的常见问题。此类场景的天空、建筑物和道路像素比行人和骑车人像素多，因为天空、建筑物和道路覆盖了图像中的更多区域。如果处理不当，这种不平衡可能对学习过程不利，因为学习会偏向于占主导的类。在此示例的稍后部分，您将使用类权重来处理此问题。

调整 CamVid 数据的大小

CamVid 数据集中的图像大小为 720×960。为减少训练时间和内存使用量，请使用 resizeCamVidImages 和 resizeCamVidPixelLabels 支持函数将图像和像素标签图像的大小调整为 360×480。

imageFolder = fullfile(outputFolder,"imagesResized",filesep);
imds = resizeCamVidImages(imds,imageFolder);

labelFolder = fullfile(outputFolder,"labelsResized",filesep);
pxds = resizeCamVidPixelLabels(pxds,labelFolder);

准备训练集和测试集

使用数据集中 60% 的图像训练 SegNet。其余图像用于测试。以下代码将图像和像素标签数据随机分成训练集和测试集。

[imdsTrain,imdsTest,pxdsTrain,pxdsTest] = partitionCamVidData(imds ...
    ,pxds);

60/40 拆分产生以下数量的训练和测试图像：

numTrainingImages = numel(imdsTrain.Files)

numTrainingImages = 421

numTestingImages = numel(imdsTest.Files)

numTestingImages = 280

创建网络

使用 fcnLayers (Computer Vision Toolbox) 创建使用 VGG-16 权重初始化的全卷积网络层。fcnLayers 函数执行网络变换以从 VGG-16 传递权重，并添加语义分割所需的其他层。fcnLayers 函数的输出是一个表示 FCN 的 LayerGraph 对象。LayerGraph 对象封装了网络层及各层之间的连接。

imageSize = [360 480];
numClasses = numel(classes);
lgraph = fcnLayers(imageSize,numClasses);

图像大小的选择基于数据集中的图像大小。类数量的选择基于 CamVid 中的类。

使用类权重平衡类

CamVid 中的类并不平衡。为了改进训练，您可以使用先前通过 countEachLabel (Computer Vision Toolbox) 函数计算的像素标签计数，计算具有中位数频率的类的权重 [3]。

imageFreq = tbl.PixelCount ./ tbl.ImagePixelCount;
classWeights = median(imageFreq) ./ imageFreq;

使用 pixelClassificationLayer (Computer Vision Toolbox) 指定类权重。

pxLayer = pixelClassificationLayer("Name","labels","Classes", ...
    tbl.Name,"ClassWeights",classWeights)

pxLayer = 
  PixelClassificationLayer with properties:

            Name: 'labels'
         Classes: [11×1 categorical]
    ClassWeights: [11×1 double]
      OutputSize: 'auto'

   Hyperparameters
    LossFunction: 'crossentropyex'

通过删除当前 pixelClassificationLayer 并添加新层，来更新具有新 pixelClassificationLayer 的 SegNet 网络。当前 pixelClassificationLayer 命名为 pixelLabels。使用 removeLayers 函数将其删除，使用 addLayers 函数添加一个新层，然后使用 connectLayers 函数将新层连接到网络的其余部分。

lgraph = removeLayers(lgraph,"pixelLabels");
lgraph = addLayers(lgraph, pxLayer);
lgraph = connectLayers(lgraph,"softmax","labels");

选择训练选项

训练的优化算法是 Adam（派生自自适应矩估计）。使用 trainingOptions 函数指定用于 Adam 的超参数。

options = trainingOptions("adam", ...
    "InitialLearnRate",1e-3, ...
    "MaxEpochs",100, ...  
    "MiniBatchSize",4, ...
    "Shuffle","every-epoch", ...
    "CheckpointPath", tempdir, ...
    "VerboseFrequency",2);

大小为 4 的 MiniBatchSize 可减少训练时的内存使用量。您可以根据系统中的 GPU 内存量增大或减小此值。

CheckpointPath 设置为临时位置。此名称-值对组让您能够在每轮训练结束时保存网络检查点。如果由于系统故障或停电而导致训练中断，您可以从保存的检查点处恢复训练。确保 CheckpointPath 指定的位置有足够的空间来存储网络检查点。

数据增强

数据增强可通过在训练期间随机变换原始数据来提高网络准确度。通过使用数据增强，您可以为训练数据添加更多变化，而不必增加带标签的训练样本的数量。要对图像和像素标签数据应用相同的随机变换，请使用数据存储组合和变换。首先，合并 imdsTrain 和 pxdsTrain。

dsTrain = combine(imdsTrain, pxdsTrain);

接下来，使用数据存储变换应用在支持函数 augmentImageAndLabel 中定义的所需数据增强。此处使用随机左/右翻转和随机 X/Y 平移 +/- 10 个像素来进行数据增强。

xTrans = [-10 10];
yTrans = [-10 10];
dsTrain = transform(dsTrain, @(data)augmentImageAndLabel(data, ...
    xTrans,yTrans));

请注意，数据增强不适用于测试数据和验证数据。理想情况下，测试数据和验证数据应代表原始数据并且保持不变，以便进行无偏置的评估。

开始训练

如果 doTraining 标志为 true，则使用 trainNetwork 开始训练。否则，请加载预训练的网络。

训练在具有 12 GB GPU 内存的 NVIDIA Titan Xp 上进行了验证。如果您的 GPU 内存较少，则可能内存不足。如果系统的内存不足，请尝试将 trainingOptions 中的 MiniBatchSize 属性降低到 1。根据您的 GPU 硬件情况，训练此网络需要大约 5 个小时或更长时间。

doTraining = false;
if doTraining    
    [net, info] = trainNetwork(dsTrain,lgraph,options);
    save("FCN8sCamVid.mat","net");
end

将 DAG 网络对象保存为名为 FCN8sCamVid.mat 的 MAT 文件。在代码生成过程中将使用此 MAT 文件。

执行 MEX 代码生成

fcn_predict 函数以图像作为输入，并使用保存在 FCN8sCamVid.mat 文件中的深度学习网络对图像执行预测。该函数将 FCN8sCamVid.mat 中的网络对象加载到持久变量 mynet 中，并在后续的预测调用中重用该持久性对象。

type("fcn_predict.m")

function out = fcn_predict(in)
%#codegen
% Copyright 2018-2019 The MathWorks, Inc.

persistent mynet;

if isempty(mynet)
    mynet = coder.loadDeepLearningNetwork('FCN8sCamVid.mat');
end

% pass in input
out = predict(mynet,in);

为 MEX 对象生成一个 GPU 配置对象以将目标语言设置为 C++。使用 coder.DeepLearningConfig (GPU Coder) 函数创建一个 cuDNN 深度学习配置对象，并将其赋给 GPU 代码配置对象的 DeepLearningConfig 属性。运行 codegen (MATLAB Coder) 命令以指定输入大小为 360×480×3。此大小对应于 FCN 的输入层。

cfg = coder.gpuConfig("mex");
cfg.TargetLang = "C++";
cfg.DeepLearningConfig = coder.DeepLearningConfig("cudnn");
codegen -config cfg fcn_predict -args {ones(360,480,3,"uint8")} -report

Code generation successful: View report

运行生成的 MEX

加载并显示输入图像。

im = imread("testImage.png");
imshow(im);

对输入图像调用 fcn_predict_mex 来运行预测。

predict_scores = fcn_predict_mex(im);

predict_scores 变量是一个三维矩阵，它具有 11 个通道，分别对应于每个类的像素级预测分数。使用最高预测分数计算通道以获得像素级标签。

[~,argmax] = max(predict_scores,[],3);

在输入图像上叠加分割标签并显示分割区域。

classes = [
    "Sky"
    "Building"
    "Pole"
    "Road"
    "Pavement"
    "Tree"
    "SignSymbol"
    "Fence"
    "Car"
    "Pedestrian"
    "Bicyclist"
    ];

cmap = camvidColorMap();
SegmentedImage = labeloverlay(im,argmax,"ColorMap",cmap);
figure
imshow(SegmentedImage);
pixelLabelColorbar(cmap,classes);

支持函数

function data = augmentImageAndLabel(data, xTrans, yTrans)
% Augment images and pixel label images using random reflection and
% translation.

for i = 1:size(data,1)
    
    tform = randomAffine2d(...
        "XReflection",true,...
        "XTranslation", xTrans, ...
        "YTranslation", yTrans);
    
    % Center the view at the center of image in the output space while
    % allowing translation to move the output image out of view.
    rout = affineOutputView(size(data{i,1}), tform, "BoundsStyle", ...
        "centerOutput");
    
    % Warp the image and pixel labels using the same transform.
    data{i,1} = imwarp(data{i,1}, tform, "OutputView", rout);
    data{i,2} = imwarp(data{i,2}, tform, "OutputView", rout);
    
end
end

参考资料

[1] Long, J., E. Shelhamer, and T. Darrell."Fully Convolutional Networks for Semantic Segmentation."Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2015, pp. 3431–3440.

[2] Brostow, G. J., J. Fauqueur, and R. Cipolla."Semantic object classes in video:A high-definition ground truth database."Pattern Recognition Letters.Vol. 30, Issue 2, 2009, pp 88-97.

[3] Badrinarayanan, V., A. Kendall, and R. Cipolla."SegNet:A Deep Convolutional Encoder-Decoder Architecture for Image Segmentation." arXiv preprint arXiv:1511.00561, 2015.