新型深度神经网络架构:ENet模型
语义分割技术能够为图像中的每个像素分配一个类别标签,这对于理解图像内容和在复杂场景中找到目标对象至关重要。在自动驾驶和增强现实等应用中,实时性是一个硬性要求,因此设计能够快速运行的卷积神经网络非常关键。
尽管深度卷积神经网络(如VGG16等)在分类和识别任务上取得了巨大成功,但它们在像素级图像标注上提供的空间结果较为粗糙。通常需要将CNN与其他算法(例如基于颜色的分割或条件随机场)级联以细化结果。许多移动或电池供电的应用需要以高于每秒10帧(fps)的速率处理图像,这对算法的运行效率提出了更高的要求。然而,现有的基于深度学习的方法往往因为参数众多和推理时间长而难以满足这一需求。
为了解决这些问题,论文提出了ENet,这是一个专为低延迟操作而设计的新型深度神经网络架构。ENet在保持相似或更高准确度的同时,显著减少了计算量、参数数量,并提高了运行速度。ENet的主要目标是在资源受限的移动设备上实现高效的语义分割,同时在高端GPU上也展现出高效的性能,以满足数据中心等场景下对大规模高分辨率图像处理的需求。
不同数据集上的ENet预测(从左到右为cityscape、CamVid和SUN):
ENet(Efficient Neural Network)
模型架构:
ENet网络框架是一种专为实时语义分割任务设计的深度学习模型,它采用了高效的编码器-解码器结构,包含多个阶段,每个阶段由多个瓶颈模块(bottleneck blocks)组成。
1. 初始阶段: 这个阶段使用一个单独的卷积块,包括一个3x3的卷积层,步长为2,用于快速将输入图像下采样,减少分辨率,同时增加特征图的数量。
2. 编码器:
瓶颈模块(Bottleneck Modules):ENet的编码器由多个瓶颈模块组成,每个模块包含以下子层:
- 1x1投影卷积:用于降维,减少特征图的数量。
- 主卷积层:可以是3x3的常规卷积、扩张卷积或5x5的不对称卷积(由5x1和1x5的卷积组成)。
- 1x1扩展卷积:用于升维,恢复特征图的数量。
- 批量归一化(Batch Normalization) 和 PReLU激活函数:在所有卷积层之后使用,以加速训练并提高模型稳定性。
- 下采样:通过最大池化层或步长为2的卷积层实现,减少特征图的空间尺寸。
3. 重复部分: 编码器中包含重复的部分,其中不包含下采样的瓶颈模块,以进一步提取特征。
4. 解码器: 解码器同样由瓶颈模块组成,但执行的是上采样操作,逐步恢复特征图的空间尺寸。
- 上采样模块:与编码器中的下采样相对应,使用上采样卷积层逐步增加特征图的分辨率。
- 最大池化与最大反池化(Max Unpooling):编码器中使用最大池化减少尺寸,在解码器中使用最大反池化恢复尺寸。
5. 全卷积层:
在解码器的最后,使用一个全卷积层将特征图转换成最终的语义分割图,该层的输出通道数等于类别数。
模型特点:
- 速度快:比现有模型快18倍。
- 计算量低:需要的浮点运算次数(FLOPs)少75倍。
- 参数少:参数数量少79倍。
- 精度高:提供与现有模型相似或更好的精度。
模型代码
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
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
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
实验步骤
数据集:在CamVid、Cityscapes和SUN数据集上进行测试。
评估指标:使用类平均精度(class average accuracy)和交并比(intersection-over-union, IoU)作为性能评估指标。
训练过程:首先训练编码器,然后添加解码器并训练整个网络进行上采样和像素级分类。
优化算法:使用Adam优化算法,设置学习率为5e-4,L2权重衰减为2e-4,批量大小为10。
实验结果
1.性能对比:ENet在NVIDIA Jetson TX1嵌入式系统模块和NVIDIA Titan X GPU上的性能均优于现有的SegNet模型。ENet在NVIDIA TX1上的推理速度能够达到21.1fps(640x360分辨率),显示出其在实时应用中的潜力:
2.精度:在Cityscapes数据集上,ENet在类IoU、类iIoU和类别IoU上均优于SegNet。在Cityscapes数据集上,ENet在类IoU上达到了58.3%,在iIoU上达到了34.4%,优于SegNet模型。
结论
ENet模型通过其创新的设计,在保持高精度的同时显著提高了语义分割的速度,使其适用于实时应用,尤其是在计算资源受限的移动设备上。尽管主要目标是在移动设备上运行网络,但在高端gpu(如NVIDIA Titan x)上也非常高效。这可能在需要处理大量高分辨率图像的数据中心应用程序中很有用。ENet允许以更快、更有效的方式执行大规模计算,这可能会大大节省成本。
- 点赞
- 收藏
- 关注作者
评论(0)