/
stupin
/
AutoDriver
Обзор
Документация
Войти
/
stupin
/
AutoDriver
Код
Запросы
0
Задачи
Вики
Пакеты
0
Релизы
0
CI/CD
Аналитика
Безопасность
master
Dev/model/ENet.py
440 строк
15 KB
Stupin
init
25 июн 2025, 19:30
25 июн 2025, 19:30
a4828f6
Код
Авторство
О чём код?
################################################################################## #ENet: A Deep Neural Network Architecture for Real-Time Semantic Segmentation #Paper-Link: https://arxiv.org/pdf/1606.02147.pdf ################################################################################## import torch import torch.nn as nn import torch.nn.functional as F from torchsummary import summary __all__ = ["ENet"] class InitialBlock(nn.Module): def __init__(self, in_channels,out_channels, kernel_size, padding=0, bias=False,relu=True): super(InitialBlock, self).__init__() if relu: activation = nn.ReLU() else: activation = nn.PReLU() self.main_branch = nn.Conv2d( in_channels, out_channels-3, kernel_size=kernel_size, stride=2, padding=padding, bias=bias, ) # MP need padding too self.ext_branch = nn.MaxPool2d(kernel_size, stride=2, padding=padding) self.batch_norm = nn.BatchNorm2d(out_channels) self.out_prelu = activation def forward(self, input): main = self.main_branch(input) ext = self.ext_branch(input) out = torch.cat((main, ext), dim=1) out = self.batch_norm(out) return self.out_prelu(out) class RegularBottleneck(nn.Module): def __init__(self, channels, internal_ratio=4, kernel_size=3, padding=0, dilation=1, asymmetric=False, dropout_prob=0., bias=False, relu=True): super(RegularBottleneck, self).__init__() internal_channels = channels // internal_ratio if relu: activation = nn.ReLU() else: activation = nn.PReLU() # 1x1 projection conv self.ext_conv1 = nn.Sequential( nn.Conv2d(channels, internal_channels, kernel_size=1, stride=1, bias=bias), nn.BatchNorm2d(internal_channels), activation, ) if asymmetric: self.ext_conv2 = nn.Sequential( nn.Conv2d(internal_channels, internal_channels, kernel_size=(kernel_size,1), stride=1, padding=(padding,0), dilation=dilation, bias=bias), nn.BatchNorm2d(internal_channels), activation, nn.Conv2d(internal_channels, internal_channels, kernel_size=(1,kernel_size), stride=1, padding=(0, padding), dilation=dilation, bias=bias), nn.BatchNorm2d(internal_channels), activation, ) else: self.ext_conv2 = nn.Sequential( nn.Conv2d(internal_channels, internal_channels, kernel_size=kernel_size, stride=1, padding=padding, dilation=dilation, bias=bias), nn.BatchNorm2d(internal_channels), activation, ) self.ext_conv3 = nn.Sequential( nn.Conv2d(internal_channels, channels, kernel_size=1, stride=1, bias=bias), nn.BatchNorm2d(channels), activation, ) self.ext_regu1 = nn.Dropout2d(p=dropout_prob) self.out_prelu = activation def forward(self, input): main = input ext = self.ext_conv1(input) ext = self.ext_conv2(ext) ext = self.ext_conv3(ext) ext = self.ext_regu1(ext) out = main + ext return self.out_prelu(out) class DownsamplingBottleneck(nn.Module): def __init__(self, in_channels, out_channels, internal_ratio=4, kernel_size=3, padding=0, return_indices=False, dropout_prob=0., bias=False, relu=True): super().__init__() # Store parameters that are needed later self.return_indices = return_indices internal_channels = in_channels // internal_ratio if relu: activation = nn.ReLU() else: activation = nn.PReLU() # Main branch - max pooling followed by feature map (channels) padding self.main_max1 = nn.MaxPool2d( kernel_size, stride=2, padding=padding, return_indices=return_indices) # Extension branch - 2x2 convolution, followed by a regular, dilated or # asymmetric convolution, followed by another 1x1 convolution. Number # of channels is doubled. # 2x2 projection convolution with stride 2, no padding self.ext_conv1 = nn.Sequential( nn.Conv2d(in_channels,internal_channels,kernel_size=2,stride=2,bias=bias), nn.BatchNorm2d(internal_channels), activation ) # Convolution self.ext_conv2 = nn.Sequential( nn.Conv2d( internal_channels, internal_channels, kernel_size=kernel_size, stride=1, padding=padding, bias=bias), nn.BatchNorm2d(internal_channels), activation) # 1x1 expansion convolution self.ext_conv3 = nn.Sequential( nn.Conv2d( internal_channels, out_channels, kernel_size=1, stride=1, bias=bias), nn.BatchNorm2d(out_channels), activation) self.ext_regul = nn.Dropout2d(p=dropout_prob) # PReLU layer to apply after concatenating the branches self.out_prelu = activation def forward(self, x): # Main branch shortcut if self.return_indices: main, max_indices = self.main_max1(x) else: main = self.main_max1(x) # Extension branch ext = self.ext_conv1(x) ext = self.ext_conv2(ext) ext = self.ext_conv3(ext) ext = self.ext_regul(ext) # Main branch channel padding # calculate for padding ch_ext - ch_main n, ch_ext, h, w = ext.size() ch_main = main.size()[1] padding = torch.zeros(n, ch_ext - ch_main, h, w) # Before concatenating, check if main is on the CPU or GPU and # convert padding accordingly if main.is_cuda: padding = padding.cuda() # Concatenate, padding for less channels of main branch main = torch.cat((main, padding), 1) # Add main and extension branches out = main + ext return self.out_prelu(out), max_indices class UpsamplingBottleneck(nn.Module): def __init__(self, in_channels, out_channels, internal_ratio=4, kernel_size=3, padding=0, dropout_prob=0., bias=False, relu=True): super().__init__() internal_channels = in_channels // internal_ratio if relu: activation = nn.ReLU() else: activation = nn.PReLU() # Main branch - max pooling followed by feature map (channels) padding self.main_conv1 = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias), nn.BatchNorm2d(out_channels)) # Remember that the stride is the same as the kernel_size, just like # the max pooling layers self.main_unpool1 = nn.MaxUnpool2d(kernel_size=2) # Extension branch - 1x1 convolution, followed by a regular, dilated or # asymmetric convolution, followed by another 1x1 convolution. Number # of channels is doubled. # 1x1 projection convolution with stride 1 self.ext_conv1 = nn.Sequential( nn.Conv2d( in_channels, internal_channels, kernel_size=1, bias=bias), nn.BatchNorm2d(internal_channels), activation) # Transposed convolution self.ext_conv2 = nn.Sequential( nn.ConvTranspose2d( internal_channels, internal_channels, kernel_size=kernel_size, stride=2, padding=padding, output_padding=1, bias=bias), nn.BatchNorm2d(internal_channels), activation) # 1x1 expansion convolution self.ext_conv3 = nn.Sequential( nn.Conv2d( internal_channels, out_channels, kernel_size=1, bias=bias), nn.BatchNorm2d(out_channels), activation) self.ext_regul = nn.Dropout2d(p=dropout_prob) # PReLU layer to apply after concatenating the branches self.out_prelu = activation def forward(self, x, max_indices): # Main branch shortcut main = self.main_conv1(x) main = self.main_unpool1(main, max_indices) # Extension branch ext = self.ext_conv1(x) ext = self.ext_conv2(ext) ext = self.ext_conv3(ext) ext = self.ext_regul(ext) # Add main and extension branches out = main + ext return self.out_prelu(out) class ENet(nn.Module): def __init__(self, classes, encoder_relu=False, decoder_relu=True): super().__init__() # source code self.name='BaseLine_ENet_trans' self.initial_block = InitialBlock(3, 16, kernel_size=3 ,padding=1, relu=encoder_relu) # Stage 1 - Encoder self.downsample1_0 = DownsamplingBottleneck( 16, 64, padding=1, return_indices=True, dropout_prob=0.01, relu=encoder_relu) self.regular1_1 = RegularBottleneck( 64, padding=1, dropout_prob=0.01, relu=encoder_relu) self.regular1_2 = RegularBottleneck( 64, padding=1, dropout_prob=0.01, relu=encoder_relu) self.regular1_3 = RegularBottleneck( 64, padding=1, dropout_prob=0.01, relu=encoder_relu) self.regular1_4 = RegularBottleneck( 64, padding=1, dropout_prob=0.01, relu=encoder_relu) # Stage 2 - Encoder self.downsample2_0 = DownsamplingBottleneck( 64, 128, padding=1, return_indices=True, dropout_prob=0.1, relu=encoder_relu) self.regular2_1 = RegularBottleneck( 128, padding=1, dropout_prob=0.1, relu=encoder_relu) self.dilated2_2 = RegularBottleneck( 128, dilation=2, padding=2, dropout_prob=0.1, relu=encoder_relu) self.asymmetric2_3 = RegularBottleneck( 128, kernel_size=5, padding=2, asymmetric=True, dropout_prob=0.1, relu=encoder_relu) self.dilated2_4 = RegularBottleneck( 128, dilation=4, padding=4, dropout_prob=0.1, relu=encoder_relu) self.regular2_5 = RegularBottleneck( 128, padding=1, dropout_prob=0.1, relu=encoder_relu) self.dilated2_6 = RegularBottleneck( 128, dilation=8, padding=8, dropout_prob=0.1, relu=encoder_relu) self.asymmetric2_7 = RegularBottleneck( 128, kernel_size=5, asymmetric=True, padding=2, dropout_prob=0.1, relu=encoder_relu) self.dilated2_8 = RegularBottleneck( 128, dilation=16, padding=16, dropout_prob=0.1, relu=encoder_relu) # Stage 3 - Encoder self.regular3_0 = RegularBottleneck( 128, padding=1, dropout_prob=0.1, relu=encoder_relu) self.dilated3_1 = RegularBottleneck( 128, dilation=2, padding=2, dropout_prob=0.1, relu=encoder_relu) self.asymmetric3_2 = RegularBottleneck( 128, kernel_size=5, padding=2, asymmetric=True, dropout_prob=0.1, relu=encoder_relu) self.dilated3_3 = RegularBottleneck( 128, dilation=4, padding=4, dropout_prob=0.1, relu=encoder_relu) self.regular3_4 = RegularBottleneck( 128, padding=1, dropout_prob=0.1, relu=encoder_relu) self.dilated3_5 = RegularBottleneck( 128, dilation=8, padding=8, dropout_prob=0.1, relu=encoder_relu) self.asymmetric3_6 = RegularBottleneck( 128, kernel_size=5, asymmetric=True, padding=2, dropout_prob=0.1, relu=encoder_relu) self.dilated3_7 = RegularBottleneck( 128, dilation=16, padding=16, dropout_prob=0.1, relu=encoder_relu) # Stage 4 - Decoder self.upsample4_0 = UpsamplingBottleneck( 128, 64, padding=1, dropout_prob=0.1, relu=decoder_relu) self.regular4_1 = RegularBottleneck( 64, padding=1, dropout_prob=0.1, relu=decoder_relu) self.regular4_2 = RegularBottleneck( 64, padding=1, dropout_prob=0.1, relu=decoder_relu) # Stage 5 - Decoder self.upsample5_0 = UpsamplingBottleneck( 64, 16, padding=1, dropout_prob=0.1, relu=decoder_relu) self.regular5_1 = RegularBottleneck( 16, padding=1, dropout_prob=0.1, relu=decoder_relu) self.transposed_conv = nn.ConvTranspose2d( 16, classes, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False) self.project_layer = nn.Conv2d(128, classes, 1, bias=False) def forward(self, x): # Initial block x = self.initial_block(x) # Stage 1 - Encoder x, max_indices1_0 = self.downsample1_0(x) x = self.regular1_1(x) x = self.regular1_2(x) x = self.regular1_3(x) x = self.regular1_4(x) # Stage 2 - Encoder x, max_indices2_0 = self.downsample2_0(x) x = self.regular2_1(x) x = self.dilated2_2(x) x = self.asymmetric2_3(x) x = self.dilated2_4(x) x = self.regular2_5(x) x = self.dilated2_6(x) x = self.asymmetric2_7(x) x = self.dilated2_8(x) # Stage 3 - Encoder x = self.regular3_0(x) x = self.dilated3_1(x) x = self.asymmetric3_2(x) x = self.dilated3_3(x) x = self.regular3_4(x) x = self.dilated3_5(x) x = self.asymmetric3_6(x) x = self.dilated3_7(x) #x = self.project_layer(x) #x = F.interpolate(x, scale_factor=8, mode='bilinear', align_corners=True) # Stage 4 - Decoder x = self.upsample4_0(x, max_indices2_0) x = self.regular4_1(x) x = self.regular4_2(x) # Stage 5 - Decoder x = self.upsample5_0(x, max_indices1_0) x = self.regular5_1(x) x = self.transposed_conv(x) return x """print layers and params of network""" if __name__ == '__main__': device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = ENet(classes=19).to(device) summary(model,(3,512,1024)) # print(model)