Vision: Making ASPP-Layer in DeepLab more generic

Created on 23 Apr 2020 · 3Comments · Source: pytorch/vision

🚀 Feature

At the moment in the ASPP-Layer defined here the number of output channels is predefined as a constant, which is good for DeepLab but not if someone else - like me ;) , want to use it in another project, where another out-channel Nr. is required.

Also the number of "atrous rates" is fixed to three, which also could be sometime more or less dpeending on the notwork. Again these fixed values may make sense in DeepLab-Model but not necessarily in other type of model, which in my case is applied on Lidar data.

that means something like this would be at least IMO better:

class ASPP(nn.Module):
    def __init__(self, in_channels, atrous_rates, out_channel=256):
        super(ASPP, self).__init__()
        modules = []
        modules.append(nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()))

        rates = tuple(atrous_rates)
        for rate in rates:
             modules.append(ASPPConv(in_channels, out_channels, rate))

        modules.append(ASPPPooling(in_channels, out_channels))

        self.convs = nn.ModuleList(modules)

        self.project = nn.Sequential(
            nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Dropout(0.5))

enhancement help wanted models semantic segmentation

Source