ENet——实时语义分割的深度神经网络架构与代码实现

概述

在移动设备上执行实时像素级分割任务具有重要意义。现有的基于分割的深度神经网络需要大量的浮点运算，并且通常需要较长时间才能投入使用。本文提出的ENet架构旨在减少潜在的计算负担。ENet在保持或提高分割精度的同时，相比现有的分割网络，速度提升了18倍，参数量减少了79倍。

论文地址：https://arxiv.org/abs/1606.02147

介绍

随着增强现实设备、智能家居设备和自动驾驶技术的兴起，将语义分割算法移植到性能较低的移动设备上变得尤为迫切。语义分割算法对图像中的每个像素进行分类标记。近期，大规模数据集的可用性以及强大的计算资源（如GPU和TPU）的出现，推动了卷积神经网络在超越传统计算机视觉算法方面取得了显著进展。尽管卷积网络在分类和识别任务上取得了良好的效果，但在进行像素级分割时，往往产生粗糙的空间结果。因此，通常需要将其他算法（如基于颜色的分割、条件随机场等）与之结合，以增强分割结果。

为了实现图像的空间分类和精细分割，已经出现了如SegNet、FCN等网络结构，这些结构通常基于VGG-16等大型多分类网络。然而，这些网络由于参数量大和推理时间长，不适用于要求图像处理速度超过10fps的移动设备或电池供电的应用设备。

本文提出的网络结构旨在实现快速推理和高准确率的分割。

相关工作

语义分割在图像理解和目标检测中扮演着关键角色，尤其在增强现实和自动驾驶中，其实时性要求极高。当前的计算机视觉应用普遍采用深度神经网络，其中场景分析较好的卷积网络采用编码-解码结构，受自编码器启发。SegNet中的编码-解码结构得到了改进，编码部分类似于VGG的卷积网络用于分类，解码部分主要用于上采样。但这些网络的参数量大，导致推理时间长。SegNet通过移除VGG16的全连接层来减少浮点运算和内存占用，但其轻量级网络仍无法实现实时分割。

其他结构采用简单的分类器后接CRF作为后处理步骤，但这些复杂的后处理操作常常无法准确标记图像中的小物体。CNN结合RNN可以提高准确率，但无法解决速度问题。RNN可作为后处理手段与其他技术结合使用。

网络结构

ENet网络结构参考了ResNet，描述为一个主分支和一个带有卷积核的附加分支，最后通过像素级相加进行融合。每个block包含三个卷积层：一个1x1的映射用于维度减少，一个主卷积层，一个1x1的扩张。在这些层之间穿插批量归一化（BN）层和PReLU层，定义为bottleneck模型。如果bottleneck进行下采样，则在主分支上添加最大池化层，并将第一个1x1的映射替换为2x2、步长为2的卷积，对激活值进行padding以匹配feature map尺寸。卷积核大小为3x3，类型包括普通卷积、空洞卷积和转置卷积，有时用1x5或5x1的非对称卷积替换。为进行正则化，在bottleneck2.0之前使用p=0.01，其他条件下使用p=0.1。

ENet的初始部分包含一个单独的block，第一阶段包含五个bottleneck部分。第二、三阶段结构相同，但第三阶段开始时没有下采样过程。前三个阶段为编码阶段，第四、五阶段为解码阶段。为减少内核调用和内存占用，网络未使用偏差项。在每个卷积层和pReLU层之间添加BN层。在解码网络部分，最大池化层被最大上采样层代替，padding被无偏差项的空间卷积替换。ENet在最后一个上采样过程中未使用最大池化索引，因为输入图像的通道数为3，而输出通道数为类别数。最终，使用全卷积模型作为网络的最后部分，占用部分解码网络的处理时间。

设计选择

Feature map尺寸：在语义分割中，对图像进行下采样存在两个缺点：一是分辨率减少导致空间位置信息损失，如边界信息；二是全像素级分割要求输出与输入尺寸相同，意味着下采样后必须进行同等程度的上采样，增加了模型大小和计算资源。为解决第一个问题，ENet采用SegNet的方式，保留最大池化过程中最大值的索引，以在解码网络中生成稀疏的上采样maps。考虑到内存需求，避免过度下采样，以免影响分割准确率。

下采样的优点：下采样后的图像进行卷积操作具有较大的感受野，有助于获取更多上下文信息，有利于不同类别的区分。研究发现使用空洞卷积效果更佳。

Early downsampling：为实现良好的分割效果和实时操作，需处理大尺寸输入图像，这非常消耗资源。ENet的前两个block大幅减小输入尺寸，同时只使用部分feature map。由于可视化信息具有高度空间冗余，可以压缩为更有效的表达形式。网络初始部分作为特征提取和预处理输入。

解码尺寸大小：SegNet具有高度对称的网络结构，而ENet是非对称的，包含较大的编码层和较小的解码网络。编码层类似于原始分类网络，处理较小数据并进行信息处理和滤波，解码网络对编码网络输出进行上采样，微调细节。

非线性操作：研究发现，去除网络初始层中的大部分ReLU层可提升分割效果，认为网络深度不足。随后，将所有ReLU替换为PReLUs，为每张feature map增加额外参数。

信息保留的维度变化：下采样是必要的，但剧烈维度衰减不利于信息流动。为解决这一问题，ENet在池化层后接卷积层以增加维度，进而增加计算资源。将池化和卷积操作并行执行，然后进行拼接。在ResNet结构中，下采样时第一个1x1映射使用步长为2的卷积，丢弃了约75%的输入信息。将卷积核增至2x2有助于信息保留。

分解卷积核：二维卷积可分解为两个一维卷积核（nxn -> nx1, 1xn）。ENet使用5x1, 1x5的非对称卷积，减小过拟合风险，增加感受野。卷积核分解还减少了参数量，加上非线性处理，使计算功能更丰富。

空洞卷积：ENet将bottleneck中的卷积层替换为空洞卷积并进行串联，增大感受野，提高分割的IOU。

正则化：现有分割数据集有限，网络训练易过拟合。ENet采用空间Dropout进行处理。

代码PyTorch实现

initial block

class InitialBlock(nn.Module):

def __init__(self,in_channels,out_channels):

super(InitialBlock, self).__init__()

self.conv = nn.Conv2d(in_channels, out_channels-in_channels, kernel_size=3, stride=2,padding=1, bias=False)

self.pool = nn.MaxPool2d(kernel_size=3,stride=2,padding=1)

self.bn = nn.BatchNorm2d(out_channels)

self.relu = nn.PReLU()

def forward(self, x):

return self.relu(self.bn(torch.cat([self.conv(x),self.pool(x)],dim=1)))

bottleneck module

class RegularBottleneck(nn.Module):

def __init__(self,in_places,places, stride=1, expansion = 4,dilation=1,is_relu=False,asymmetric=False,p=0.01):

super(RegularBottleneck, self).__init__()

mid_channels = in_places // expansion        self.bottleneck = nn.Sequential(

Conv1x1BNReLU(in_places, mid_channels, False),

AsymmetricConv(mid_channels, 1, is_relu) if asymmetric else Conv3x3BNReLU(mid_channels, mid_channels, 1,dilation, is_relu),

Conv1x1BNReLU(mid_channels, places,is_relu),

nn.Dropout2d(p=p)

)

self.relu = nn.ReLU(inplace=True) if is_relu else nn.PReLU()

def forward(self, x):

residual = x        out = self.bottleneck(x)

out += residual        out = self.relu(out)

return outclass DownBottleneck(nn.Module):

def __init__(self,in_places,places, stride=2, expansion = 4,is_relu=False,p=0.01):

super(DownBottleneck, self).__init__()

mid_channels = in_places // expansion        self.bottleneck = nn.Sequential(

Conv2x2BNReLU(in_places, mid_channels, is_relu),

Conv3x3BNReLU(mid_channels, mid_channels, 1, 1, is_relu),

Conv1x1BNReLU(mid_channels, places,is_relu),

nn.Dropout2d(p=p)

)

self.downsample = nn.MaxPool2d(3,stride=stride,padding=1,return_indices=True)

self.relu = nn.ReLU(inplace=True) if is_relu else nn.PReLU()

def forward(self, x):

out = self.bottleneck(x)

residual,indices = self.downsample(x)

n, ch, h, w = out.size()

ch_res = residual.size()[1]

padding = torch.zeros(n, ch - ch_res, h, w)

residual = torch.cat((residual, padding), 1)

out += residual        out = self.relu(out)

return out, indicesclass UpBottleneck(nn.Module):

def __init__(self,in_places,places, stride=2, expansion = 4,is_relu=True,p=0.01):

super(UpBottleneck, self).__init__()

mid_channels = in_places // expansion        self.bottleneck = nn.Sequential(

Conv1x1BNReLU(in_places,mid_channels,is_relu),

TransposeConv3x3BNReLU(mid_channels,mid_channels,stride,is_relu),

Conv1x1BNReLU(mid_channels,places,is_relu),

nn.Dropout2d(p=p)

)

self.upsample_conv = Conv1x1BN(in_places, places)

self.upsample_unpool = nn.MaxUnpool2d(kernel_size=2)

self.relu = nn.ReLU(inplace=True) if is_relu else nn.PReLU()

def forward(self, x, indices):

out = self.bottleneck(x)

residual = self.upsample_conv(x)

residual = self.upsample_unpool(residual,indices)

out += residual        out = self.relu(out)

return

architecture

class ENet(nn.Module):

def __init__(self, num_classes):

super(ENet, self).__init__()

self.initialBlock = InitialBlock(3,16)

self.stage1_1 = DownBottleneck(16, 64, 2)

self.stage1_2 = nn.Sequential(

RegularBottleneck(64, 64, 1),

RegularBottleneck(64, 64, 1),

RegularBottleneck(64, 64, 1),

RegularBottleneck(64, 64, 1),

)

self.stage2_1 = DownBottleneck(64, 128, 2)

self.stage2_2 = nn.Sequential(

RegularBottleneck(128, 128, 1),

RegularBottleneck(128, 128, 1, dilation=2),

RegularBottleneck(128, 128, 1, asymmetric=True),

RegularBottleneck(128, 128, 1, dilation=4),

RegularBottleneck(128, 128, 1),

RegularBottleneck(128, 128, 1, dilation=8),

RegularBottleneck(128, 128, 1, asymmetric=True),

RegularBottleneck(128, 128, 1, dilation=16),

)

self.stage3 = nn.Sequential(

RegularBottleneck(128, 128, 1),

RegularBottleneck(128, 128, 1, dilation=2),

RegularBottleneck(128, 128, 1, asymmetric=True),

RegularBottleneck(128, 128, 1, dilation=4),

RegularBottleneck(128, 128, 1),

RegularBottleneck(128, 128, 1, dilation=8),

RegularBottleneck(128, 128, 1, asymmetric=True),

RegularBottleneck(128, 128, 1, dilation=16),

)

self.stage4_1 = UpBottleneck(128, 64, 2, is_relu=True)

self.stage4_2 = nn.Sequential(

RegularBottleneck(64, 64, 1, is_relu=True),

RegularBottleneck(64, 64, 1, is_relu=True),

)

self.stage5_1 = UpBottleneck(64, 16, 2, is_relu=True)

self.stage5_2 = RegularBottleneck(16, 16, 1, is_relu=True)

self.final_conv = nn.ConvTranspose2d(in_channels=16, out_channels=num_classes, kernel_size=3, stride=2, padding=1,

output_padding=1, bias=False)

def forward(self, x):

x = self.initialBlock(x)

x,indices1 = self.stage1_1(x)

x = self.stage1_2(x)

x, indices2 = self.stage2_1(x)

x = self.stage2_2(x)

x = self.stage3(x)

x = self.stage4_1(x, indices2)

x = self.stage4_2(x)

x = self.stage5_1(x, indices1)

x = self.stage5_2(x)

out = self.final_conv(x)

return