Spaces:
Running
Running
| import layers | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| from . import spec_utils | |
| class BaseASPPNet(nn.Module): | |
| def __init__(self, nin, ch, dilations=(4, 8, 16)): | |
| super(BaseASPPNet, self).__init__() | |
| self.enc1 = layers.Encoder(nin, ch, 3, 2, 1) | |
| self.enc2 = layers.Encoder(ch, ch * 2, 3, 2, 1) | |
| self.enc3 = layers.Encoder(ch * 2, ch * 4, 3, 2, 1) | |
| self.enc4 = layers.Encoder(ch * 4, ch * 8, 3, 2, 1) | |
| self.aspp = layers.ASPPModule(ch * 8, ch * 16, dilations) | |
| self.dec4 = layers.Decoder(ch * (8 + 16), ch * 8, 3, 1, 1) | |
| self.dec3 = layers.Decoder(ch * (4 + 8), ch * 4, 3, 1, 1) | |
| self.dec2 = layers.Decoder(ch * (2 + 4), ch * 2, 3, 1, 1) | |
| self.dec1 = layers.Decoder(ch * (1 + 2), ch, 3, 1, 1) | |
| def __call__(self, x): | |
| h, e1 = self.enc1(x) | |
| h, e2 = self.enc2(h) | |
| h, e3 = self.enc3(h) | |
| h, e4 = self.enc4(h) | |
| h = self.aspp(h) | |
| h = self.dec4(h, e4) | |
| h = self.dec3(h, e3) | |
| h = self.dec2(h, e2) | |
| h = self.dec1(h, e1) | |
| return h | |
| class CascadedASPPNet(nn.Module): | |
| def __init__(self, n_fft): | |
| super(CascadedASPPNet, self).__init__() | |
| self.stg1_low_band_net = BaseASPPNet(2, 16) | |
| self.stg1_high_band_net = BaseASPPNet(2, 16) | |
| self.stg2_bridge = layers.Conv2DBNActiv(18, 8, 1, 1, 0) | |
| self.stg2_full_band_net = BaseASPPNet(8, 16) | |
| self.stg3_bridge = layers.Conv2DBNActiv(34, 16, 1, 1, 0) | |
| self.stg3_full_band_net = BaseASPPNet(16, 32) | |
| self.out = nn.Conv2d(32, 2, 1, bias=False) | |
| self.aux1_out = nn.Conv2d(16, 2, 1, bias=False) | |
| self.aux2_out = nn.Conv2d(16, 2, 1, bias=False) | |
| self.max_bin = n_fft // 2 | |
| self.output_bin = n_fft // 2 + 1 | |
| self.offset = 128 | |
| def forward(self, x, aggressiveness=None): | |
| mix = x.detach() | |
| x = x.clone() | |
| x = x[:, :, : self.max_bin] | |
| bandw = x.size()[2] // 2 | |
| aux1 = torch.cat( | |
| [ | |
| self.stg1_low_band_net(x[:, :, :bandw]), | |
| self.stg1_high_band_net(x[:, :, bandw:]), | |
| ], | |
| dim=2, | |
| ) | |
| h = torch.cat([x, aux1], dim=1) | |
| aux2 = self.stg2_full_band_net(self.stg2_bridge(h)) | |
| h = torch.cat([x, aux1, aux2], dim=1) | |
| h = self.stg3_full_band_net(self.stg3_bridge(h)) | |
| mask = torch.sigmoid(self.out(h)) | |
| mask = F.pad( | |
| input=mask, | |
| pad=(0, 0, 0, self.output_bin - mask.size()[2]), | |
| mode="replicate", | |
| ) | |
| if self.training: | |
| aux1 = torch.sigmoid(self.aux1_out(aux1)) | |
| aux1 = F.pad( | |
| input=aux1, | |
| pad=(0, 0, 0, self.output_bin - aux1.size()[2]), | |
| mode="replicate", | |
| ) | |
| aux2 = torch.sigmoid(self.aux2_out(aux2)) | |
| aux2 = F.pad( | |
| input=aux2, | |
| pad=(0, 0, 0, self.output_bin - aux2.size()[2]), | |
| mode="replicate", | |
| ) | |
| return mask * mix, aux1 * mix, aux2 * mix | |
| else: | |
| if aggressiveness: | |
| mask[:, :, : aggressiveness["split_bin"]] = torch.pow( | |
| mask[:, :, : aggressiveness["split_bin"]], | |
| 1 + aggressiveness["value"] / 3, | |
| ) | |
| mask[:, :, aggressiveness["split_bin"] :] = torch.pow( | |
| mask[:, :, aggressiveness["split_bin"] :], | |
| 1 + aggressiveness["value"], | |
| ) | |
| return mask * mix | |
| def predict(self, x_mag, aggressiveness=None): | |
| h = self.forward(x_mag, aggressiveness) | |
| if self.offset > 0: | |
| h = h[:, :, :, self.offset : -self.offset] | |
| assert h.size()[3] > 0 | |
| return h | |