| 220 | return out |
| 221 | |
| 222 | class DFNet(nn.Module): |
| 223 | def __init__(self, width=64, input_channel_rate=1, output_channel=2): |
| 224 | |
| 225 | super(DFNet, self).__init__() |
| 226 | # self.device = device |
| 227 | self.in_channels = 2 * input_channel_rate |
| 228 | self.out_channels = output_channel |
| 229 | self.kernel_size = (2, 3) |
| 230 | # self.elu = nn.SELU(inplace=True) |
| 231 | self.pad = nn.ConstantPad2d((1, 1, 1, 0), value=0.) |
| 232 | self.pad1 = nn.ConstantPad2d((1, 1, 0, 0), value=0.) |
| 233 | self.width = width |
| 234 | |
| 235 | self.inp_conv = nn.Conv2d(in_channels=self.in_channels, out_channels=self.width, kernel_size=(1, 1)) # [b, 64, nframes, 256] |
| 236 | self.inp_norm = InstantLayerNorm2d(width) |
| 237 | self.inp_prelu = nn.PReLU(self.width) |
| 238 | |
| 239 | self.enc_dense1 = DenseBlock(4, self.width) |
| 240 | self.dual_transformer = Dual_Transformer(self.width, self.width, num_layers=4) # # [b, 64, nframes, 8] |
| 241 | |
| 242 | # gated output layer |
| 243 | self.output1 = nn.Sequential( |
| 244 | nn.Conv2d(in_channels=self.width, out_channels=self.width, kernel_size=1), |
| 245 | nn.Tanh() |
| 246 | ) |
| 247 | self.output2 = nn.Sequential( |
| 248 | nn.Conv2d(in_channels=self.width, out_channels=self.width, kernel_size=1), |
| 249 | nn.Sigmoid() |
| 250 | ) |
| 251 | |
| 252 | self.dec_dense1 = DenseBlock(4, self.width) |
| 253 | |
| 254 | self.out_conv = nn.Conv2d(in_channels=self.width, out_channels=self.out_channels, kernel_size=(1, 1)) |
| 255 | |
| 256 | |
| 257 | def forward(self, x): |
| 258 | |
| 259 | x = x.permute(0,1,3,2) # [B, 2, num_frames, num_bins] |
| 260 | out = self.inp_prelu(self.inp_norm(self.inp_conv(x))) # [b, 64, num_frames, frame_size] |
| 261 | out = self.enc_dense1(out) # [b, 64, num_frames, frame_size] |
| 262 | |
| 263 | out = self.dual_transformer(out) # [b, 64, num_frames, 256] |
| 264 | out = self.output1(out) * self.output2(out) # mask [b, 64, num_frames, 256] |
| 265 | out = self.dec_dense1(out) |
| 266 | out = self.out_conv(out) |
| 267 | out = out.permute(0,1,3,2) |
| 268 | return out |