| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| from . import spec_utils |
|
|
|
|
| class Conv2DBNActiv(nn.Module): |
| def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU): |
| super(Conv2DBNActiv, self).__init__() |
| self.conv = nn.Sequential( |
| nn.Conv2d( |
| nin, |
| nout, |
| kernel_size=ksize, |
| stride=stride, |
| padding=pad, |
| dilation=dilation, |
| bias=False, |
| ), |
| nn.BatchNorm2d(nout), |
| activ(), |
| ) |
|
|
| def __call__(self, x): |
| return self.conv(x) |
|
|
|
|
| class SeperableConv2DBNActiv(nn.Module): |
| def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU): |
| super(SeperableConv2DBNActiv, self).__init__() |
| self.conv = nn.Sequential( |
| nn.Conv2d( |
| nin, |
| nin, |
| kernel_size=ksize, |
| stride=stride, |
| padding=pad, |
| dilation=dilation, |
| groups=nin, |
| bias=False, |
| ), |
| nn.Conv2d(nin, nout, kernel_size=1, bias=False), |
| nn.BatchNorm2d(nout), |
| activ(), |
| ) |
|
|
| def __call__(self, x): |
| return self.conv(x) |
|
|
|
|
| class Encoder(nn.Module): |
| def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU): |
| super(Encoder, self).__init__() |
| self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ) |
| self.conv2 = Conv2DBNActiv(nout, nout, ksize, stride, pad, activ=activ) |
|
|
| def __call__(self, x): |
| skip = self.conv1(x) |
| h = self.conv2(skip) |
|
|
| return h, skip |
|
|
|
|
| class Decoder(nn.Module): |
| def __init__( |
| self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False |
| ): |
| super(Decoder, self).__init__() |
| self.conv = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ) |
| self.dropout = nn.Dropout2d(0.1) if dropout else None |
|
|
| def __call__(self, x, skip=None): |
| x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True) |
| if skip is not None: |
| skip = spec_utils.crop_center(skip, x) |
| x = torch.cat([x, skip], dim=1) |
| h = self.conv(x) |
|
|
| if self.dropout is not None: |
| h = self.dropout(h) |
|
|
| return h |
|
|
|
|
| class ASPPModule(nn.Module): |
| def __init__(self, nin, nout, dilations=(4, 8, 16), activ=nn.ReLU): |
| super(ASPPModule, self).__init__() |
| self.conv1 = nn.Sequential( |
| nn.AdaptiveAvgPool2d((1, None)), |
| Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ), |
| ) |
| self.conv2 = Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ) |
| self.conv3 = SeperableConv2DBNActiv( |
| nin, nin, 3, 1, dilations[0], dilations[0], activ=activ |
| ) |
| self.conv4 = SeperableConv2DBNActiv( |
| nin, nin, 3, 1, dilations[1], dilations[1], activ=activ |
| ) |
| self.conv5 = SeperableConv2DBNActiv( |
| nin, nin, 3, 1, dilations[2], dilations[2], activ=activ |
| ) |
| self.bottleneck = nn.Sequential( |
| Conv2DBNActiv(nin * 5, nout, 1, 1, 0, activ=activ), nn.Dropout2d(0.1) |
| ) |
|
|
| def forward(self, x): |
| _, _, h, w = x.size() |
| feat1 = F.interpolate( |
| self.conv1(x), size=(h, w), mode="bilinear", align_corners=True |
| ) |
| feat2 = self.conv2(x) |
| feat3 = self.conv3(x) |
| feat4 = self.conv4(x) |
| feat5 = self.conv5(x) |
| out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1) |
| bottle = self.bottleneck(out) |
| return bottle |
|
|
