ENet: A Deep Neural Network Architecture for Real-Time Semantic Segmentation
ENet: A Deep Neural Network Architecture for Real-Time Semantic Segmentation
논문 출처 : https://arxiv.org/pdf/1606.02147.pdf
초록
실시간 시멘틱 세그먼테이션을 수행하는 능력은 모바일 어플리케이션에서 매우 중요하다. 긴 floating point 동작의 불리하게하며 긴 실행시간을 가지게한다. 논문에서 제안한 ENET은 낮은 지연시간을 요하는 작업을 위해 만들어졌다. E-Net은 기존 대비 18배 이상 빠르며 75배 보다 적은 FLOPS, 79배 적은 파라미터 값을 가지고, 기존 모델에 비해 더 나은 정확도를 보이고 있다. CamVid, Cityscapes, Sun 데이터셋에서 현재 최고의 모델 (segnet)과 비교해본다. 그리고 정확도와 프로세싱 시간의 트레이트 오프 관계를 비교한다. 그리고 임베디드 시스템(TX1)에서 성능 측정과 ENET을 더 빠르게할 가능한 SW 개선점을 제시한다.
서론
웨어러블, 홈 자동화 장치, 자율주행 자동차는 낮은 파워에서 실시간 동작의 시멘틱 세그먼테이션이 필요하다. SegNet이나 fully convolution networkssms VGG16 구조를 기반으로 한다. 이런 큰 모델은 엄청나게 많은 파라미터와 긴 inference 시간을 갖는다. 10fps 이상이 요구되는 이미지 처리가 필요한 핸드폰이나 낮은 배터리 파워 어플리케이션에서는 사용할 수 없다. 이 논문에서는 높은 정확도와 빠른 inference를 위해 최적화된 새로운 네트워크 제안한다. ENET이라고 빠르고 콤팩트한 엔코더 디코더 구조이다.
결과
ENET을 사용해서 Camvid, Cityscape, SUN RGB-D 를 벤치마크했다. SegNet을 기준선으로 잡았는데 이유는 가장 빠른 segmentaion 모델이고 적은 파라미터와 FCN보다 작은 메모리 사용 때문이다.
성능 분석
추론시간
하드웨어 요구사항
소프트웨어 한계 이러한 성능 수준에 도달할 수 있었던 가장 중요한 기술 중 하나는 convolutional layer factorization 입니다. 하지만 우리는 한 가지 놀라운 결점을 발견했습니다. 이 방법을 적용하면 부동 소수점 연산 및 파라미터 수를 크게 줄일 수 있었지만 개별 커널 호출 수도 증가하여 각각 크기가 작아졌습니다.
이러한 작업 중 일부는 너무 저렴해져서 GPU 커널 론칭 비용이 실제 컴퓨팅 비용을 초과하기 시작한다는 사실을 알게 되었습니다. 또 기존 커널이 레지스터에 보관된 값에 접근할 수 없어 출시 시 글로벌 메모리에서 모든 데이터를 로드하고 작업이 끝나면 저장해야 합니다. 즉, 기능 맵을 지속적으로 저장하고 다시 로드해야 하기 때문에 더 많은 수의 커널을 사용하면 메모리 트랜잭션 수가 증가합니다. 이는 비선형 연산의 경우 특히 명백해집니다. ENet에서 PReLU는 추론 시간의 4분의 1 이상을 소비합니다. 단순한 점별 연산에 불과하고 병렬화가 매우 쉽기 때문에 앞에서 언급한 데이터 이동으로 인해 발생한다고 가정합니다. 이는 심각한 제한 사항이지만, 기존 소프트웨어에서 커널 퓨전을 수행함으로써 해결할 수 있습니다. 즉, 변환 결과에 비선형성을 직접 적용하는 커널을 생성하거나 한 번의 호출로 여러 개의 작은 변환을 수행하면 해결할 수 있습니다. cuDNN과 같은 GPU 라이브러리의 이러한 향상은 네트워크의 속도와 효율성을 더욱 높일 수 있습니다.
벤치마크
Cityscapes
CamVid
SUN RGB-D
네크워크 구조
- 1x1 conv를 통한 차원 축소
main conv(regular, dilated, full convolution(a.k.a deconvolution)
- 1x1 conv를 통한 차원 확대
- Downsampling인 경우, max pooling 연산이 main branch에 추가되며, 1x1 projection 컨볼루션을 대신하여 stride 2의 2x2 컨볼루션을 대체하고, zero pad 활성화를 수행한다.
- Asymmetric인 경우, 5x5 컨볼루션을 5x1과 1x5 컨볼루션의 연속으로 대체하였다.
- Regularizer는 spatial dropout을 사용하는데, bottleneck 2.0 이전에는 p=0.01, 이후에는 0.1로 설정하였다.
Encoder에서 Stage 1은 5개의 병목 블록으로 구성되어 있고, 2단계와 3단계도 동일한 구조를 가지고 있다. 단 3단계는 처음에 입력을 다운샘플링 하지 않는다. Stage 3 단계까지 인코더이며, 4단계와 5단계는 디코더에 해당한다. Decoder에서는 max pooling은 max unpooling으로 대체되고, padding은 bias가 없는 spatial convolution으로 대체된다.
최종 출력은 C 특징맵이다. ( 오브젝트 클래스의 개수이다. )
Design Choices
Feature map resolution Down-sampling에는 두가지 주요 단점이 있다. 첫째, 특징맵 해상도를 줄이면 가장자리 부분의 공간정보를 손실하게 된다. 둘째, 전체 픽셀 분할은 출력이 입력과 동일한 해상도를 가져야 한다.
첫째 문제는 인코더에 의해 생성된 피처맵을 추가하여 FCN에서와 같이, max-pooling layer에서 선택된 요소의 인덱스를 저장하고 이를 사용하여 디코더에서 sparse 업 샘플링 된 맵을 생성하는 방법을 사용하였다. 여전히, 정확도가 낮아지는 문제가 있지만 가능한 한 제한적인 영향을 주도록하였다. 그러나, 다운 샘플링 연산은 더 큰 receptive field를 가지기 때문에 더 많은 context 정보를 얻을 수 있다. 예를 들어, 도로 장면에서 라이더와 보행자 같은 클래스를 구분하려 할 때, 사람의 외모만을 학습하는 것만큼이나 맥락을 학습하는 것도 중요하기 때문이다.
Early downsampling real-time 연산에서 가장 중요한 것은 입력 프레임에 대해 효율적으로 처리하는 것이며, E-net의 처음 두 블록은 입력 크기를 크게 줄이고 작은 특징 맵 세트만 사용합니다.
Decoder size SegNet에서와 같이, 대칭적인 구조를 가지며 인코더의 경우 더 작은 해상도 데이터에서 작동하고 정보 처리 및 필터링을 제공 할 수 있어야 한다. 반면, 디코더의 역할은 인코더의 출력을 업샘플링하고 세부 사항만 미세 조정하게 된다.
Nonlinear operations 최근 ReLU와 batch normalization을 사용하는 것이 유익하다고 보고되어 있으며, 네트워크 초기 계층에서 대부분의 ReLU를 제거하면 결과가 향상된다는 사실을 발견했다. 또한, 음의 기울기를 학습하기 위해, ReLU 대신에 PReLU를 사용하여 기울기를 학습하도록 하였다. PReLU의 가중치 분포는 그림 3과 같다.
Information-preserving dimensionality changes
Factorizing filters n x n 컨볼루션 연산을 n x 1과 1xn 으로 나누어 수행할 수 있으며, 중복연산을 피할 수 있다. 또한, 병목 모듈에서 사용되는 연산은 하나의 큰 컨볼루션 레이어를 낮은 순위 근사치인 더 작고 간단한 작업으로 분해하는 것으로 볼 수 있다. 이러한 분해는 속도를 크게 높이고 매개 변수의 수를 크게 줄여 중복성을 줄일 수 있다. 또한 레이어 사이에 삽입되는 비 선형 연산 덕분에 계산 기능을 더 풍부하게 만들 수 있다.
Dilated convolution 넓은 수용 필드를 갖는 것이 매우 중요하며, 더 넓은 Context를 고려하여 분류를 수행할 수 잇다. 특징 맵을 과도하게 다운 샘플링하지 않고 개선하기 위해 확장된 convolution을 사용하기로 결정하였다. 가장 작은 해상도에서 작동하는 컨볼루션 연산을 대체하였으며, 추가 비용없이 Cityscapes의 IoU를 4% 높여 정확도를 크게 향상 시켰다.
Regularization 대부분의 픽셀 단위 분할 데이터 세트는 상대적으로 작기 때문에 신경망과 같은 표현모델이 빠르게 과적합되기 시작한다. 초기 실험에서 L2 weight decay를 사용했고, 확률적 깊이(stochastic)를 시도하여 정확도를 높였다. 우리는 컨볼루션 브랜치 끝에 spatial dropout을 배치했는데, 확률적 깊이보다 훨씬 더 작동하는 것으로 나타났다.
코드
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from torch.optim.lr_scheduler import StepLR
import cv2
import os
from tqdm import tqdm
from PIL import Image
class InitialBlock(nn.Module):
# Initial block of the model:
# Input
# / \\\\
# / \\\\
#maxpool2d conv2d-3x3
# \\\\ /
# \\\\ /
# concatenate
def __init__ (self,in_channels = 3,out_channels = 13):
super().__init__()
self.maxpool = nn.MaxPool2d(kernel_size=2,stride = 2,padding = 0)
self.conv = nn.Conv2d(in_channels,out_channels,kernel_size = 3,stride = 2,padding = 1)
self.prelu = nn.PReLU(16)
self.batchnorm = nn.BatchNorm2d(out_channels)
def forward(self, x):
main = self.conv(x)
main = self.batchnorm(main)
side = self.maxpool(x)
# concatenating on the channels axis
x = torch.cat((main, side), dim=1)
x = self.prelu(x)
return x
class UBNeck(nn.Module):
# Upsampling bottleneck:
# Bottleneck Input
# / \\\\
# / \\\\
# conv2d-1x1 convTrans2d-1x1
# | | PReLU
# | convTrans2d-3x3
# | | PReLU
# | convTrans2d-1x1
# | |
# maxunpool2d Regularizer
# \\\\ /
# \\\\ /
# Summing + PReLU
#
# Params:
# projection_ratio - ratio between input and output channels
# relu - if True: relu used as the activation function else: Prelu us used
def __init__(self, in_channels, out_channels, relu=False, projection_ratio=4):
super().__init__()
# Define class variables
self.in_channels = in_channels
self.reduced_depth = int(in_channels / projection_ratio)
self.out_channels = out_channels
if relu:
activation = nn.ReLU()
else:
activation = nn.PReLU()
self.unpool = nn.MaxUnpool2d(kernel_size = 2,stride = 2)
self.main_conv = nn.Conv2d(in_channels = self.in_channels,out_channels = self.out_channels,kernel_size = 1)
self.dropout = nn.Dropout2d(p=0.1)
self.convt1 = nn.ConvTranspose2d(in_channels = self.in_channels,out_channels = self.reduced_depth,kernel_size = 1,padding = 0,bias = False)
self.prelu1 = activation
# This layer used for Upsampling
self.convt2 = nn.ConvTranspose2d(in_channels = self.reduced_depth,out_channels = self.reduced_depth,kernel_size = 3,stride = 2,padding = 1,output_padding = 1,bias = False)
self.prelu2 = activation
self.convt3 = nn.ConvTranspose2d(in_channels = self.reduced_depth,out_channels = self.out_channels,kernel_size = 1,padding = 0,bias = False)
self.prelu3 = activation
self.batchnorm = nn.BatchNorm2d(self.reduced_depth)
self.batchnorm2 = nn.BatchNorm2d(self.out_channels)
def forward(self, x, indices):
x_copy = x
# Side Branch
x = self.convt1(x)
x = self.batchnorm(x)
x = self.prelu1(x)
x = self.convt2(x)
x = self.batchnorm(x)
x = self.prelu2(x)
x = self.convt3(x)
x = self.batchnorm2(x)
x = self.dropout(x)
# Main Branch
x_copy = self.main_conv(x_copy)
x_copy = self.unpool(x_copy, indices, output_size=x.size())
# summing the main and side branches
x = x + x_copy
x = self.prelu3(x)
return x
class RDDNeck(nn.Module):
def __init__(self, dilation, in_channels, out_channels, down_flag, relu=False, projection_ratio=4, p=0.1):
# Regular|Dilated|Downsampling bottlenecks:
#
# Bottleneck Input
# / \\\\
# / \\\\
# maxpooling2d conv2d-1x1
# | | PReLU
# | conv2d-3x3
# | | PReLU
# | conv2d-1x1
# | |
# Padding2d Regularizer
# \\\\ /
# \\\\ /
# Summing + PReLU
#
# Params:
# dilation (bool) - if True: creating dilation bottleneck
# down_flag (bool) - if True: creating downsampling bottleneck
# projection_ratio - ratio between input and output channels
# relu - if True: relu used as the activation function else: Prelu us used
# p - dropout ratio
super().__init__()
# Define class variables
self.in_channels = in_channels
self.out_channels = out_channels
self.dilation = dilation
self.down_flag = down_flag
# calculating the number of reduced channels
if down_flag:
self.stride = 2
self.reduced_depth = int(in_channels // projection_ratio)
else:
self.stride = 1
self.reduced_depth = int(out_channels // projection_ratio)
if relu:
activation = nn.ReLU()
else:
activation = nn.PReLU()
self.maxpool = nn.MaxPool2d(kernel_size = 2,
stride = 2,
padding = 0, return_indices=True)
self.dropout = nn.Dropout2d(p=p)
self.conv1 = nn.Conv2d(in_channels = self.in_channels,out_channels = self.reduced_depth,kernel_size = 1,stride = 1,padding = 0,bias = False,dilation = 1)
self.prelu1 = activation
self.conv2 = nn.Conv2d(in_channels = self.reduced_depth,out_channels = self.reduced_depth,kernel_size = 3,stride = self.stride,padding = self.dilation,bias = True,dilation = self.dilation)
self.prelu2 = activation
self.conv3 = nn.Conv2d(in_channels = self.reduced_depth,out_channels = self.out_channels,kernel_size = 1,stride = 1,padding = 0,bias = False,dilation = 1)
self.prelu3 = activation
self.batchnorm = nn.BatchNorm2d(self.reduced_depth)
self.batchnorm2 = nn.BatchNorm2d(self.out_channels)
def forward(self, x):
bs = x.size()[0]
x_copy = x
# Side Branch
x = self.conv1(x)
x = self.batchnorm(x)
x = self.prelu1(x)
x = self.conv2(x)
x = self.batchnorm(x)
x = self.prelu2(x)
x = self.conv3(x)
x = self.batchnorm2(x)
x = self.dropout(x)
# Main Branch
if self.down_flag:
x_copy, indices = self.maxpool(x_copy)
if self.in_channels != self.out_channels:
out_shape = self.out_channels - self.in_channels
#padding and concatenating in order to match the channels axis of the side and main branches
extras = torch.zeros((bs, out_shape, x.shape[2], x.shape[3]))
if torch.cuda.is_available():
extras = extras.cuda()
x_copy = torch.cat((x_copy, extras), dim = 1)
# Summing main and side branches
x = x + x_copy
x = self.prelu3(x)
if self.down_flag:
return x, indices
else:
return x
class ASNeck(nn.Module):
def __init__(self, in_channels, out_channels, projection_ratio=4):
# Asymetric bottleneck:
#
# Bottleneck Input
# / \\\\
# / \\\\
# | conv2d-1x1
# | | PReLU
# | conv2d-1x5
# | |
# | conv2d-5x1
# | | PReLU
# | conv2d-1x1
# | |
# Padding2d Regularizer
# \\\\ /
# \\\\ /
# Summing + PReLU
#
# Params:
# projection_ratio - ratio between input and output channels
super().__init__()
# Define class variables
self.in_channels = in_channels
self.reduced_depth = int(in_channels / projection_ratio)
self.out_channels = out_channels
self.dropout = nn.Dropout2d(p=0.1)
self.conv1 = nn.Conv2d(in_channels = self.in_channels,out_channels = self.reduced_depth,kernel_size = 1,stride = 1,adding = 0,bias = False)
self.prelu1 = nn.PReLU()
self.conv21 = nn.Conv2d(in_channels = self.reduced_depth,out_channels = self.reduced_depth,kernel_size = (1, 5),stride = 1,padding = (0, 2),bias = False)
self.conv22 = nn.Conv2d(in_channels = self.reduced_depth,out_channels = self.reduced_depth,kernel_size = (5, 1),stride = 1,padding = (2, 0),bias = False)
self.prelu2 = nn.PReLU()
self.conv3 = nn.Conv2d(in_channels = self.reduced_depth,out_channels = self.out_channels,kernel_size = 1,stride = 1,padding = 0,bias = False)
self.prelu3 = nn.PReLU()
self.batchnorm = nn.BatchNorm2d(self.reduced_depth)
self.batchnorm2 = nn.BatchNorm2d(self.out_channels)
def forward(self, x):
bs = x.size()[0]
x_copy = x
# Side Branch
x = self.conv1(x)
x = self.batchnorm(x)
x = self.prelu1(x)
x = self.conv21(x)
x = self.conv22(x)
x = self.batchnorm(x)
x = self.prelu2(x)
x = self.conv3(x)
x = self.dropout(x)
x = self.batchnorm2(x)
# Main Branch
if self.in_channels != self.out_channels:
out_shape = self.out_channels - self.in_channels
#padding and concatenating in order to match the channels axis of the side and main branches
extras = torch.zeros((bs, out_shape, x.shape[2], x.shape[3]))
if torch.cuda.is_available():
extras = extras.cuda()
x_copy = torch.cat((x_copy, extras), dim = 1)
# Summing main and side branches
x = x + x_copy
x = self.prelu3(x)
return x
class ENet(nn.Module):
# Creating Enet model!
def __init__(self, C):
super().__init__()
# Define class variables
# C - number of classes
self.C = C
# The initial block
self.init = InitialBlock()
# The first bottleneck
self.b10 = RDDNeck(dilation=1,in_channels=16, out_channels=64,down_flag=True,p=0.01)
self.b11 = RDDNeck(dilation=1,in_channels=64,out_channels=64,down_flag=False,p=0.01)
self.b12 = RDDNeck(dilation=1,in_channels=64,out_channels=64,down_flag=False,p=0.01)
self.b13 = RDDNeck(dilation=1,in_channels=64,out_channels=64,down_flag=False,p=0.01)
self.b14 = RDDNeck(dilation=1,in_channels=64,out_channels=64,down_flag=False,p=0.01)
# The second bottleneck
self.b20 = RDDNeck(dilation=1,in_channels=64,out_channels=128,down_flag=True)
self.b21 = RDDNeck(dilation=1,in_channels=128,out_channels=128,own_flag=False)
self.b22 = RDDNeck(dilation=2,in_channels=128,out_channels=128,down_flag=False)
self.b23 = ASNeck(in_channels=128,out_channels=128)
self.b24 = RDDNeck(dilation=4,in_channels=128,out_channels=128,down_flag=False)
self.b25 = RDDNeck(dilation=1,in_channels=128,out_channels=128,down_flag=False)
self.b26 = RDDNeck(dilation=8,in_channels=128,out_channels=128,down_flag=False)
self.b27 = ASNeck(in_channels=128,out_channels=128)
self.b28 = RDDNeck(dilation=16,in_channels=128,out_channels=128,down_flag=False)
# The third bottleneck
self.b31 = RDDNeck(dilation=1,in_channels=128,out_channels=128,down_flag=False)
self.b32 = RDDNeck(dilation=2,in_channels=128,out_channels=128,down_flag=False)
self.b33 = ASNeck(in_channels=128,out_channels=128)
self.b34 = RDDNeck(dilation=4,in_channels=128,out_channels=128,down_flag=False)
self.b35 = RDDNeck(dilation=1,in_channels=128,out_channels=128,down_flag=False)
self.b36 = RDDNeck(dilation=8,in_channels=128,out_channels=128,down_flag=False)
self.b37 = ASNeck(in_channels=128,out_channels=128)
self.b38 = RDDNeck(dilation=16,in_channels=128,out_channels=128,down_flag=False)
# The fourth bottleneck
self.b40 = UBNeck(in_channels=128,out_channels=64,relu=True)
self.b41 = RDDNeck(dilation=1,in_channels=64,out_channels=64,down_flag=False,relu=True)
self.b42 = RDDNeck(dilation=1,in_channels=64,out_channels=64,down_flag=False,relu=True)
# The fifth bottleneck
self.b50 = UBNeck(in_channels=64,out_channels=16,relu=True)
self.b51 = RDDNeck(dilation=1,in_channels=16,out_channels=16,down_flag=False,relu=True)
# Final ConvTranspose Layer
self.fullconv = nn.ConvTranspose2d(in_channels=16,out_channels=self.C,kernel_size=3,stride=2,padding=1,output_padding=1,bias=False)
def forward(self, x):
# The initial block
x = self.init(x)
# The first bottleneck
x, i1 = self.b10(x)
x = self.b11(x)
x = self.b12(x)
x = self.b13(x)
x = self.b14(x)
# The second bottleneck
x, i2 = self.b20(x)
x = self.b21(x)
x = self.b22(x)
x = self.b23(x)
x = self.b24(x)
x = self.b25(x)
x = self.b26(x)
x = self.b27(x)
x = self.b28(x)
# The third bottleneck
x = self.b31(x)
x = self.b32(x)
x = self.b33(x)
x = self.b34(x)
x = self.b35(x)
x = self.b36(x)
x = self.b37(x)
x = self.b38(x)
# The fourth bottleneck
x = self.b40(x, i2)
x = self.b41(x)
x = self.b42(x)
# The fifth bottleneck
x = self.b50(x, i1)
x = self.b51(x)
# Final ConvTranspose Layer
x = self.fullconv(x)
return x