"""LEDNet: A Lightweight Encoder-Decoder Network for Real-time Semantic Segmentation""" import torch import torch.nn as nn import torch.nn.functional as F from core.nn import _ConvBNReLU __all__ = ['LEDNet', 'get_lednet', 'get_lednet_citys'] class LEDNet(nn.Module): r"""LEDNet Parameters ---------- nclass : int Number of categories for the training dataset. backbone : string Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', 'resnet101' or 'resnet152'). norm_layer : object Normalization layer used in backbone network (default: :class:`nn.BatchNorm`; for Synchronized Cross-GPU BachNormalization). aux : bool Auxiliary loss. Reference: Yu Wang, et al. "LEDNet: A Lightweight Encoder-Decoder Network for Real-Time Semantic Segmentation." arXiv preprint arXiv:1905.02423 (2019). """ def __init__(self, nclass, backbone='', aux=False, jpu=False, pretrained_base=True, **kwargs): super(LEDNet, self).__init__() self.encoder = nn.Sequential( Downsampling(3, 32), SSnbt(32, **kwargs), SSnbt(32, **kwargs), SSnbt(32, **kwargs), Downsampling(32, 64), SSnbt(64, **kwargs), SSnbt(64, **kwargs), Downsampling(64, 128), SSnbt(128, **kwargs), SSnbt(128, 2, **kwargs), SSnbt(128, 5, **kwargs), SSnbt(128, 9, **kwargs), SSnbt(128, 2, **kwargs), SSnbt(128, 5, **kwargs), SSnbt(128, 9, **kwargs), SSnbt(128, 17, **kwargs), ) self.decoder = APNModule(128, nclass) self.__setattr__('exclusive', ['encoder', 'decoder']) def forward(self, x): size = x.size()[2:] x = self.encoder(x) x = self.decoder(x) outputs = list() x = F.interpolate(x, size, mode='bilinear', align_corners=True) outputs.append(x) #return tuple(outputs) return outputs[0] class Downsampling(nn.Module): def __init__(self, in_channels, out_channels, **kwargs): super(Downsampling, self).__init__() self.conv1 = nn.Conv2d(in_channels, out_channels // 2, 3, 2, 2, bias=False) self.conv2 = nn.Conv2d(in_channels, out_channels // 2, 3, 2, 2, bias=False) self.pool = nn.MaxPool2d(kernel_size=2, stride=1) def forward(self, x): x1 = self.conv1(x) x1 = self.pool(x1) x2 = self.conv2(x) x2 = self.pool(x2) return torch.cat([x1, x2], dim=1) class SSnbt(nn.Module): def __init__(self, in_channels, dilation=1, norm_layer=nn.BatchNorm2d, **kwargs): super(SSnbt, self).__init__() inter_channels = in_channels // 2 self.branch1 = nn.Sequential( nn.Conv2d(inter_channels, inter_channels, (3, 1), padding=(1, 0), bias=False), nn.ReLU(True), nn.Conv2d(inter_channels, inter_channels, (1, 3), padding=(0, 1), bias=False), norm_layer(inter_channels), nn.ReLU(True), nn.Conv2d(inter_channels, inter_channels, (3, 1), padding=(dilation, 0), dilation=(dilation, 1), bias=False), nn.ReLU(True), nn.Conv2d(inter_channels, inter_channels, (1, 3), padding=(0, dilation), dilation=(1, dilation), bias=False), norm_layer(inter_channels), nn.ReLU(True)) self.branch2 = nn.Sequential( nn.Conv2d(inter_channels, inter_channels, (1, 3), padding=(0, 1), bias=False), nn.ReLU(True), nn.Conv2d(inter_channels, inter_channels, (3, 1), padding=(1, 0), bias=False), norm_layer(inter_channels), nn.ReLU(True), nn.Conv2d(inter_channels, inter_channels, (1, 3), padding=(0, dilation), dilation=(1, dilation), bias=False), nn.ReLU(True), nn.Conv2d(inter_channels, inter_channels, (3, 1), padding=(dilation, 0), dilation=(dilation, 1), bias=False), norm_layer(inter_channels), nn.ReLU(True)) self.relu = nn.ReLU(True) @staticmethod def channel_shuffle(x, groups): n, c, h, w = x.size() channels_per_group = c // groups x = x.view(n, groups, channels_per_group, h, w) x = torch.transpose(x, 1, 2).contiguous() x = x.view(n, -1, h, w) return x def forward(self, x): # channels split x1, x2 = x.split(x.size(1) // 2, 1) x1 = self.branch1(x1) x2 = self.branch2(x2) out = torch.cat([x1, x2], dim=1) out = self.relu(out + x) out = self.channel_shuffle(out, groups=2) return out class APNModule(nn.Module): def __init__(self, in_channels, nclass, norm_layer=nn.BatchNorm2d, **kwargs): super(APNModule, self).__init__() self.conv1 = _ConvBNReLU(in_channels, in_channels, 3, 2, 1, norm_layer=norm_layer) self.conv2 = _ConvBNReLU(in_channels, in_channels, 5, 2, 2, norm_layer=norm_layer) self.conv3 = _ConvBNReLU(in_channels, in_channels, 7, 2, 3, norm_layer=norm_layer) self.level1 = _ConvBNReLU(in_channels, nclass, 1, norm_layer=norm_layer) self.level2 = _ConvBNReLU(in_channels, nclass, 1, norm_layer=norm_layer) self.level3 = _ConvBNReLU(in_channels, nclass, 1, norm_layer=norm_layer) self.level4 = _ConvBNReLU(in_channels, nclass, 1, norm_layer=norm_layer) self.level5 = nn.Sequential( nn.AdaptiveAvgPool2d(1), _ConvBNReLU(in_channels, nclass, 1)) def forward(self, x): w, h = x.size()[2:] branch3 = self.conv1(x) branch2 = self.conv2(branch3) branch1 = self.conv3(branch2) out = self.level1(branch1) out = F.interpolate(out, ((w + 3) // 4, (h + 3) // 4), mode='bilinear', align_corners=True) out = self.level2(branch2) + out out = F.interpolate(out, ((w + 1) // 2, (h + 1) // 2), mode='bilinear', align_corners=True) out = self.level3(branch3) + out out = F.interpolate(out, (w, h), mode='bilinear', align_corners=True) out = self.level4(x) * out out = self.level5(x) + out return out def get_lednet(dataset='citys', backbone='', pretrained=False, root='~/.torch/models', pretrained_base=True, **kwargs): acronyms = { 'pascal_voc': 'pascal_voc', 'pascal_aug': 'pascal_aug', 'ade20k': 'ade', 'coco': 'coco', 'citys': 'citys', } from ..data.dataloader import datasets model = LEDNet(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs) if pretrained: from .model_store import get_model_file device = torch.device(kwargs['local_rank']) model.load_state_dict(torch.load(get_model_file('lednet_%s' % (acronyms[dataset]), root=root), map_location=device)) return model def get_lednet_citys(**kwargs): return get_lednet('citys', **kwargs) if __name__ == '__main__': #model = get_lednet_citys() input = torch.rand(2, 3, 224, 224) model =LEDNet(4, pretrained_base=True) # target = torch.zeros(4, 512, 512).cuda() # model.eval() # print(model) loss = model(input) print(loss, loss.shape) # from torchsummary import summary # # summary(model, (3, 224, 224)) # 打印表格，按顺序输出每层的输出形状和参数 import torch from thop import profile from torchsummary import summary flop, params = profile(model, input_size=(1, 3, 512, 512)) print('flops:{:.3f}G\nparams:{:.3f}M'.format(flop / 1e9, params / 1e6))