| 237 | |
| 238 | |
| 239 | class StreamEncoder(nn.Module): |
| 240 | def __init__(self): |
| 241 | super().__init__() |
| 242 | self.en_convs = nn.ModuleList([ |
| 243 | ConvBlock(3*3, 16, (1,5), stride=(1,2), padding=(0,2), use_deconv=False, is_last=False), |
| 244 | ConvBlock(16, 16, (1,5), stride=(1,2), padding=(0,2), groups=2, use_deconv=False, is_last=False), |
| 245 | StreamGTConvBlock(16, 16, (3,3), stride=(1,1), padding=(0,1), dilation=(1,1), use_deconv=False), |
| 246 | StreamGTConvBlock(16, 16, (3,3), stride=(1,1), padding=(0,1), dilation=(2,1), use_deconv=False), |
| 247 | StreamGTConvBlock(16, 16, (3,3), stride=(1,1), padding=(0,1), dilation=(5,1), use_deconv=False) |
| 248 | ]) |
| 249 | |
| 250 | def forward(self, x, conv_cache, tra_cache): |
| 251 | """ |
| 252 | x: (B,C,T,F) |
| 253 | conv_cache: (B,C, (kT-1)*8, F) |
| 254 | tra_cache: (3,1,B,C) |
| 255 | """ |
| 256 | en_outs = [] |
| 257 | for i in range(2): |
| 258 | x = self.en_convs[i](x) |
| 259 | en_outs.append(x) |
| 260 | |
| 261 | x, conv_cache[:,:, :2, :], tra_cache[0] = self.en_convs[2](x, conv_cache[:,:, :2, :], tra_cache[0]); en_outs.append(x) |
| 262 | x, conv_cache[:,:, 2:6, :], tra_cache[1] = self.en_convs[3](x, conv_cache[:,:, 2:6, :], tra_cache[1]); en_outs.append(x) |
| 263 | x, conv_cache[:,:, 6:16, :], tra_cache[2] = self.en_convs[4](x, conv_cache[:,:, 6:16, :], tra_cache[2]); en_outs.append(x) |
| 264 | |
| 265 | return x, en_outs, conv_cache, tra_cache |
| 266 | |
| 267 | |
| 268 | class StreamDecoder(nn.Module): |