'''网络结构定义'''



from __future__ import absolute_import, print_function 
from collections import OrderedDict 
import torch 
import torch.nn as nn 
import torch.nn.functional as F 
class _ImagePool(nn.Module): 
    def __init__(self, in_ch, out_ch): 
        super().__init__() 
        self.pool = nn.AdaptiveAvgPool2d(1) 
        self.conv = _ConvBnReLU(in_ch, out_ch, 1, 1, 0, 1) 
    def forward(self, x):
        _, _, H, W = x.shape 
        h = self.pool(x) 
        h = self.conv(h) 
        h = F.interpolate(h, size=(H, W), mode="bilinear", align_corners=False) 
        return h 
        
class _ASPP(nn.Module):
    """ Atrous spatial pyramid pooling with image-level feature """ 
    def __init__(self, in_ch, out_ch, rates): 
        super(_ASPP, self).__init__() 
        self.stages = nn.Module() 
        self.stages.add_module("c0", _ConvBnReLU(in_ch, out_ch, 1, 1, 0, 1)) 
        for i, rate in enumerate(rates): 
            self.stages.add_module( "c{}".format(i + 1), _ConvBnReLU(in_ch, out_ch, 3, 1, padding=rate, dilation=rate), ) 
            self.stages.add_module("imagepool", _ImagePool(in_ch, out_ch)) 
                
    def forward(self, x): 
        return torch.cat([stage(x) for stage in self.stages.children()], dim=1) 

class DeepLabV3(nn.Sequential): 
    """ DeepLab v3: Dilated ResNet with multi-grid + improved ASPP """ 
    def __init__(self, n_classes, n_blocks, atrous_rates, multi_grids, output_stride): 
        super(DeepLabV3, self).__init__() 
        # Stride and dilation 
        if output_stride == 8: 
            s = [1, 2, 1, 1] 
            d = [1, 1, 2, 4] 
        elif output_stride == 16: 
            s = [1, 2, 2, 1] 
            d = [1, 1, 1, 2] 
        ch = [64 * 2 ** p for p in range(6)] 
        self.add_module("layer1", _Stem(ch[0])) 
        self.add_module("layer2", _ResLayer(n_blocks[0], ch[0], ch[2], s[0], d[0])) 
        self.add_module("layer3", _ResLayer(n_blocks[1], ch[2], ch[3], s[1], d[1])) 
        self.add_module("layer4", _ResLayer(n_blocks[2], ch[3], ch[4], s[2], d[2])) 
        self.add_module( "layer5", _ResLayer(n_blocks[3], ch[4], ch[5], s[3], d[3], multi_grids) ) 
        self.add_module("aspp", _ASPP(ch[5], 256, atrous_rates)) 
        concat_ch = 256 * (len(atrous_rates) + 2) 
        self.add_module("fc1", _ConvBnReLU(concat_ch, 256, 1, 1, 0, 1)) 
        self.add_module("fc2", nn.Conv2d(256, n_classes, kernel_size=1)) 
try: 
    from encoding.nn import SyncBatchNorm 
    _BATCH_NORM = SyncBatchNorm 
except: 
    _BATCH_NORM = nn.BatchNorm2d 
_BOTTLENECK_EXPANSION = 4 

class _ConvBnReLU(nn.Sequential): 
    """ Cascade of 2D convolution, batch norm, and ReLU. """ 
    BATCH_NORM = _BATCH_NORM 
    def __init__( self, in_ch, out_ch, kernel_size, stride, padding, dilation, relu=True ): 
        super(_ConvBnReLU, self).__init__() 
        self.add_module( "conv", nn.Conv2d( in_ch, out_ch, kernel_size, stride, padding, dilation, bias=False ), ) 
        self.add_module("bn", _BATCH_NORM(out_ch, eps=1e-5, momentum=0.999)) 
        if relu: 
            self.add_module("relu", nn.ReLU()) 

class _Bottleneck(nn.Module): 
    """ Bottleneck block of MSRA ResNet. """ 
    def __init__(self, in_ch, out_ch, stride, dilation, downsample): 
        super(_Bottleneck, self).__init__() 
        mid_ch = out_ch // _BOTTLENECK_EXPANSION 
        self.reduce = _ConvBnReLU(in_ch, mid_ch, 1, stride, 0, 1, True) 
        self.conv3x3 = _ConvBnReLU(mid_ch, mid_ch, 3, 1, dilation, dilation, True) 
        self.increase = _ConvBnReLU(mid_ch, out_ch, 1, 1, 0, 1, False) 
        self.shortcut = ( 
        _ConvBnReLU(in_ch, out_ch, 1, stride, 0, 1, False) 
        if downsample 
        else lambda x: x # identity 
        ) 
    
    def forward(self, x): 
        h = self.reduce(x) 
        h = self.conv3x3(h) 
        h = self.increase(h) 
        h += self.shortcut(x) 
        return F.relu(h) 
        
class _ResLayer(nn.Sequential): 
    """ Residual layer with multi grids """ 
    def __init__(self, n_layers, in_ch, out_ch, stride, dilation, multi_grids=None): 
        super(_ResLayer, self).__init__() 
        if multi_grids is None: 
            multi_grids = [1 for _ in range(n_layers)] 
        else: assert n_layers == len(multi_grids) # Downsampling is only in the first block 
        for i in range(n_layers): 
            self.add_module( 
                "block{}".format(i + 1), 
                _Bottleneck( 
                    in_ch=(in_ch if i == 0 else out_ch), 
                    out_ch=out_ch, 
                    stride=(stride if i == 0 else 1), 
                    dilation=dilation * multi_grids[i], 
                    downsample=(True if i == 0 else False), 
                ), 
            ) 
            
class _Stem(nn.Sequential): 
    """ The 1st conv layer. Note that the max pooling is different from both MSRA and FAIR ResNet. """ 
    def __init__(self, out_ch): 
        super(_Stem, self).__init__() 
        self.add_module("conv1", _ConvBnReLU(3, out_ch, 7, 2, 3, 1)) 
        self.add_module("pool", nn.MaxPool2d(3, 2, 1, ceil_mode=True)) 

if __name__ == "__main__": 
    model = DeepLabV3( n_classes=21, n_blocks=[3, 4, 23, 3], atrous_rates=[6, 12, 18], multi_grids=[1, 2, 4], output_stride=8, ) 
    model.eval() 
    image = torch.randn(1, 3, 513, 513) 
    print(model) 
    print("input:", image.shape) 
    print("output:", model(image).shape)
'''网络结构定义'''



from __future__ import absolute_import, print_function 
from collections import OrderedDict 
import torch 
import torch.nn as nn 
import torch.nn.functional as F 
class _ImagePool(nn.Module): 
    def __init__(self, in_ch, out_ch): 
        super().__init__() 
        self.pool = nn.AdaptiveAvgPool2d(1) 
        self.conv = _ConvBnReLU(in_ch, out_ch, 1, 1, 0, 1) 
    def forward(self, x):
        _, _, H, W = x.shape 
        h = self.pool(x) 
        h = self.conv(h) 
        h = F.interpolate(h, size=(H, W), mode="bilinear", align_corners=False) 
        return h 
        
class _ASPP(nn.Module):
    """ Atrous spatial pyramid pooling with image-level feature """ 
    def __init__(self, in_ch, out_ch, rates): 
        super(_ASPP, self).__init__() 
        self.stages = nn.Module() 
        self.stages.add_module("c0", _ConvBnReLU(in_ch, out_ch, 1, 1, 0, 1)) 
        for i, rate in enumerate(rates): 
            self.stages.add_module( "c{}".format(i + 1), _ConvBnReLU(in_ch, out_ch, 3, 1, padding=rate, dilation=rate), ) 
            self.stages.add_module("imagepool", _ImagePool(in_ch, out_ch)) 
                
    def forward(self, x): 
        return torch.cat([stage(x) for stage in self.stages.children()], dim=1) 

class DeepLabV3(nn.Sequential): 
    """ DeepLab v3: Dilated ResNet with multi-grid + improved ASPP """ 
    def __init__(self, n_classes, n_blocks, atrous_rates, multi_grids, output_stride): 
        super(DeepLabV3, self).__init__() 
        # Stride and dilation 
        if output_stride == 8: 
            s = [1, 2, 1, 1] 
            d = [1, 1, 2, 4] 
        elif output_stride == 16: 
            s = [1, 2, 2, 1] 
            d = [1, 1, 1, 2] 
        ch = [64 * 2 ** p for p in range(6)] 
        self.add_module("layer1", _Stem(ch[0])) 
        self.add_module("layer2", _ResLayer(n_blocks[0], ch[0], ch[2], s[0], d[0])) 
        self.add_module("layer3", _ResLayer(n_blocks[1], ch[2], ch[3], s[1], d[1])) 
        self.add_module("layer4", _ResLayer(n_blocks[2], ch[3], ch[4], s[2], d[2])) 
        self.add_module( "layer5", _ResLayer(n_blocks[3], ch[4], ch[5], s[3], d[3], multi_grids) ) 
        self.add_module("aspp", _ASPP(ch[5], 256, atrous_rates)) 
        concat_ch = 256 * (len(atrous_rates) + 2) 
        self.add_module("fc1", _ConvBnReLU(concat_ch, 256, 1, 1, 0, 1)) 
        self.add_module("fc2", nn.Conv2d(256, n_classes, kernel_size=1)) 
try: 
    from encoding.nn import SyncBatchNorm 
    _BATCH_NORM = SyncBatchNorm 
except: 
    _BATCH_NORM = nn.BatchNorm2d 
_BOTTLENECK_EXPANSION = 4 

class _ConvBnReLU(nn.Sequential): 
    """ Cascade of 2D convolution, batch norm, and ReLU. """ 
    BATCH_NORM = _BATCH_NORM 
    def __init__( self, in_ch, out_ch, kernel_size, stride, padding, dilation, relu=True ): 
        super(_ConvBnReLU, self).__init__() 
        self.add_module( "conv", nn.Conv2d( in_ch, out_ch, kernel_size, stride, padding, dilation, bias=False ), ) 
        self.add_module("bn", _BATCH_NORM(out_ch, eps=1e-5, momentum=0.999)) 
        if relu: 
            self.add_module("relu", nn.ReLU()) 

class _Bottleneck(nn.Module): 
    """ Bottleneck block of MSRA ResNet. """ 
    def __init__(self, in_ch, out_ch, stride, dilation, downsample): 
        super(_Bottleneck, self).__init__() 
        mid_ch = out_ch // _BOTTLENECK_EXPANSION 
        self.reduce = _ConvBnReLU(in_ch, mid_ch, 1, stride, 0, 1, True) 
        self.conv3x3 = _ConvBnReLU(mid_ch, mid_ch, 3, 1, dilation, dilation, True) 
        self.increase = _ConvBnReLU(mid_ch, out_ch, 1, 1, 0, 1, False) 
        self.shortcut = ( 
        _ConvBnReLU(in_ch, out_ch, 1, stride, 0, 1, False) 
        if downsample 
        else lambda x: x # identity 
        ) 
    
    def forward(self, x): 
        h = self.reduce(x) 
        h = self.conv3x3(h) 
        h = self.increase(h) 
        h += self.shortcut(x) 
        return F.relu(h) 
        
class _ResLayer(nn.Sequential): 
    """ Residual layer with multi grids """ 
    def __init__(self, n_layers, in_ch, out_ch, stride, dilation, multi_grids=None): 
        super(_ResLayer, self).__init__() 
        if multi_grids is None: 
            multi_grids = [1 for _ in range(n_layers)] 
        else: assert n_layers == len(multi_grids) # Downsampling is only in the first block 
        for i in range(n_layers): 
            self.add_module( 
                "block{}".format(i + 1), 
                _Bottleneck( 
                    in_ch=(in_ch if i == 0 else out_ch), 
                    out_ch=out_ch, 
                    stride=(stride if i == 0 else 1), 
                    dilation=dilation * multi_grids[i], 
                    downsample=(True if i == 0 else False), 
                ), 
            ) 
            
class _Stem(nn.Sequential): 
    """ The 1st conv layer. Note that the max pooling is different from both MSRA and FAIR ResNet. """ 
    def __init__(self, out_ch): 
        super(_Stem, self).__init__() 
        self.add_module("conv1", _ConvBnReLU(3, out_ch, 7, 2, 3, 1)) 
        self.add_module("pool", nn.MaxPool2d(3, 2, 1, ceil_mode=True)) 

if __name__ == "__main__": 
    model = DeepLabV3( n_classes=21, n_blocks=[3, 4, 23, 3], atrous_rates=[6, 12, 18], multi_grids=[1, 2, 4], output_stride=8, ) 
    model.eval() 
    image = torch.randn(1, 3, 513, 513) 
    print(model) 
    print("input:", image.shape) 
    print("output:", model(image).shape)
网络结构定义 Atrous spatial pyramid pooling with image-level feature DeepLab v3: Dilated ResNet with multi-grid + improved ASPP # Stride and dilation Cascade of 2D convolution, batch norm, and ReLU. Bottleneck block of MSRA ResNet. # identity Residual layer with multi grids # Downsampling is only in the first block The 1st conv layer. Note that the max pooling is different from both MSRA and FAIR ResNet.