| 266 | |
| 267 | |
| 268 | class StreamDecoder(nn.Module): |
| 269 | def __init__(self): |
| 270 | super().__init__() |
| 271 | self.de_convs = nn.ModuleList([ |
| 272 | StreamGTConvBlock(16, 16, (3,3), stride=(1,1), padding=(0,1), dilation=(5,1), use_deconv=True), |
| 273 | StreamGTConvBlock(16, 16, (3,3), stride=(1,1), padding=(0,1), dilation=(2,1), use_deconv=True), |
| 274 | StreamGTConvBlock(16, 16, (3,3), stride=(1,1), padding=(0,1), dilation=(1,1), use_deconv=True), |
| 275 | ConvBlock(16, 16, (1,5), stride=(1,2), padding=(0,2), groups=2, use_deconv=True, is_last=False), |
| 276 | ConvBlock(16, 2, (1,5), stride=(1,2), padding=(0,2), use_deconv=True, is_last=True) |
| 277 | ]) |
| 278 | |
| 279 | def forward(self, x, en_outs, conv_cache, tra_cache): |
| 280 | """ |
| 281 | x: (B,C,T,F) |
| 282 | conv_cache: (B,C, (kT-1)*8, F) |
| 283 | tra_cache: (3,1,B,C) |
| 284 | """ |
| 285 | x, conv_cache[:,:, 6:16, :], tra_cache[0] = self.de_convs[0](x + en_outs[4], conv_cache[:,:, 6:16, :], tra_cache[0]) |
| 286 | x, conv_cache[:,:, 2:6, :], tra_cache[1] = self.de_convs[1](x + en_outs[3], conv_cache[:,:, 2:6, :], tra_cache[1]) |
| 287 | x, conv_cache[:,:, :2, :], tra_cache[2] = self.de_convs[2](x + en_outs[2], conv_cache[:,:, :2, :], tra_cache[2]) |
| 288 | |
| 289 | for i in range(3, 5): |
| 290 | x = self.de_convs[i](x + en_outs[4-i]) |
| 291 | return x, conv_cache, tra_cache |
| 292 | |
| 293 | |
| 294 | class Mask(nn.Module): |