| 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 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, stride, pad, activ=activ)
|
| self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
|
|
|
| def __call__(self, x):
|
| h = self.conv1(x)
|
| h = self.conv2(h)
|
|
|
| return h
|
|
|
|
|
| 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.conv1 = 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.conv1(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, 12), activ=nn.ReLU, dropout=False):
|
| super(ASPPModule, self).__init__()
|
| self.conv1 = nn.Sequential(
|
| nn.AdaptiveAvgPool2d((1, None)),
|
| Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ),
|
| )
|
| self.conv2 = Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ)
|
| self.conv3 = Conv2DBNActiv(
|
| nin, nout, 3, 1, dilations[0], dilations[0], activ=activ
|
| )
|
| self.conv4 = Conv2DBNActiv(
|
| nin, nout, 3, 1, dilations[1], dilations[1], activ=activ
|
| )
|
| self.conv5 = Conv2DBNActiv(
|
| nin, nout, 3, 1, dilations[2], dilations[2], activ=activ
|
| )
|
| self.bottleneck = Conv2DBNActiv(nout * 5, nout, 1, 1, 0, activ=activ)
|
| self.dropout = nn.Dropout2d(0.1) if dropout else None
|
|
|
| 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)
|
| out = self.bottleneck(out)
|
|
|
| if self.dropout is not None:
|
| out = self.dropout(out)
|
|
|
| return out
|
|
|
|
|
| class LSTMModule(nn.Module):
|
| def __init__(self, nin_conv, nin_lstm, nout_lstm):
|
| super(LSTMModule, self).__init__()
|
| self.conv = Conv2DBNActiv(nin_conv, 1, 1, 1, 0)
|
| self.lstm = nn.LSTM(
|
| input_size=nin_lstm, hidden_size=nout_lstm // 2, bidirectional=True
|
| )
|
| self.dense = nn.Sequential(
|
| nn.Linear(nout_lstm, nin_lstm), nn.BatchNorm1d(nin_lstm), nn.ReLU()
|
| )
|
|
|
| def forward(self, x):
|
| N, _, nbins, nframes = x.size()
|
| h = self.conv(x)[:, 0]
|
| h = h.permute(2, 0, 1)
|
| h, _ = self.lstm(h)
|
| h = self.dense(h.reshape(-1, h.size()[-1]))
|
| h = h.reshape(nframes, N, 1, nbins)
|
| h = h.permute(1, 2, 3, 0)
|
|
|
| return h
|
|
|
