| 245 | |
| 246 | |
| 247 | class Decoder(nn.Module): |
| 248 | def __init__(self): |
| 249 | super().__init__() |
| 250 | self.de_convs = nn.ModuleList([ |
| 251 | GTConvBlock(16, 16, (3,3), stride=(1,1), padding=(2*5,1), dilation=(5,1), use_deconv=True), |
| 252 | GTConvBlock(16, 16, (3,3), stride=(1,1), padding=(2*2,1), dilation=(2,1), use_deconv=True), |
| 253 | GTConvBlock(16, 16, (3,3), stride=(1,1), padding=(2*1,1), dilation=(1,1), use_deconv=True), |
| 254 | ConvBlock(16, 16, (1,5), stride=(1,2), padding=(0,2), groups=2, use_deconv=True, is_last=False), |
| 255 | ConvBlock(16, 2, (1,5), stride=(1,2), padding=(0,2), use_deconv=True, is_last=True) |
| 256 | ]) |
| 257 | |
| 258 | def forward(self, x, en_outs): |
| 259 | N_layers = len(self.de_convs) |
| 260 | for i in range(N_layers): |
| 261 | x = self.de_convs[i](x + en_outs[N_layers-1-i]) |
| 262 | return x |
| 263 | |
| 264 | |
| 265 | class Mask(nn.Module): |