jone/Music_Source_Separation
3
1import math2from typing import List3 4import numpy as np5import matplotlib.pyplot as plt6import pytorch_lightning as pl7import torch8import torch.nn as nn9import torch.nn.functional as F10import torch.optim as optim11from torch.optim.lr_scheduler import LambdaLR12from torchlibrosa.stft import STFT, ISTFT, magphase13 14from bytesep.models.pytorch_modules import (15 Base,16 init_bn,17 init_embedding,18 init_layer,19 act,20 Subband,21)22 23 24class ConvBlock(nn.Module):25 def __init__(26 self,27 in_channels,28 out_channels,29 condition_size,30 kernel_size,31 activation,32 momentum,33 ):34 super(ConvBlock, self).__init__()35 36 self.activation = activation37 padding = (kernel_size[0] // 2, kernel_size[1] // 2)38 39 self.conv1 = nn.Conv2d(40 in_channels=in_channels,41 out_channels=out_channels,42 kernel_size=kernel_size,43 stride=(1, 1),44 dilation=(1, 1),45 padding=padding,46 bias=False,47 )48 49 self.bn1 = nn.BatchNorm2d(out_channels, momentum=momentum)50 51 self.conv2 = nn.Conv2d(52 in_channels=out_channels,53 out_channels=out_channels,54 kernel_size=kernel_size,55 stride=(1, 1),56 dilation=(1, 1),57 padding=padding,58 bias=False,59 )60 61 self.bn2 = nn.BatchNorm2d(out_channels, momentum=momentum)62 63 self.beta1 = nn.Linear(condition_size, out_channels, bias=True)64 self.beta2 = nn.Linear(condition_size, out_channels, bias=True)65 66 self.init_weights()67 68 def init_weights(self):69 init_layer(self.conv1)70 init_layer(self.conv2)71 init_bn(self.bn1)72 init_bn(self.bn2)73 init_embedding(self.beta1)74 init_embedding(self.beta2)75 76 def forward(self, x, condition):77 78 b1 = self.beta1(condition)[:, :, None, None]79 b2 = self.beta2(condition)[:, :, None, None]80 81 x = act(self.bn1(self.conv1(x)) + b1, self.activation)82 x = act(self.bn2(self.conv2(x)) + b2, self.activation)83 return x84 85 86class EncoderBlock(nn.Module):87 def __init__(88 self,89 in_channels,90 out_channels,91 condition_size,92 kernel_size,93 downsample,94 activation,95 momentum,96 ):97 super(EncoderBlock, self).__init__()98 99 self.conv_block = ConvBlock(100 in_channels, out_channels, condition_size, kernel_size, activation, momentum101 )102 self.downsample = downsample103 104 def forward(self, x, condition):105 encoder = self.conv_block(x, condition)106 encoder_pool = F.avg_pool2d(encoder, kernel_size=self.downsample)107 return encoder_pool, encoder108 109 110class DecoderBlock(nn.Module):111 def __init__(112 self,113 in_channels,114 out_channels,115 condition_size,116 kernel_size,117 upsample,118 activation,119 momentum,120 ):121 super(DecoderBlock, self).__init__()122 self.kernel_size = kernel_size123 self.stride = upsample124 self.activation = activation125 126 self.conv1 = torch.nn.ConvTranspose2d(127 in_channels=in_channels,128 out_channels=out_channels,129 kernel_size=self.stride,130 stride=self.stride,131 padding=(0, 0),132 bias=False,133 dilation=(1, 1),134 )135 136 self.bn1 = nn.BatchNorm2d(out_channels, momentum=momentum)137 138 self.conv_block2 = ConvBlock(139 out_channels * 2,140 out_channels,141 condition_size,142 kernel_size,143 activation,144 momentum,145 )146 147 self.beta1 = nn.Linear(condition_size, out_channels, bias=True)148 149 self.init_weights()150 151 def init_weights(self):152 init_layer(self.conv1)153 init_bn(self.bn1)154 init_embedding(self.beta1)155 156 def forward(self, input_tensor, concat_tensor, condition):157 b1 = self.beta1(condition)[:, :, None, None]158 x = act(self.bn1(self.conv1(input_tensor)) + b1, self.activation)159 x = torch.cat((x, concat_tensor), dim=1)160 x = self.conv_block2(x, condition)161 return x162 163 164class ConditionalUNet(nn.Module, Base):165 def __init__(self, input_channels, target_sources_num):166 super(ConditionalUNet, self).__init__()167 168 self.input_channels = input_channels169 condition_size = target_sources_num170 self.output_sources_num = 1171 172 window_size = 2048173 hop_size = 441174 center = True175 pad_mode = "reflect"176 window = "hann"177 activation = "relu"178 momentum = 0.01179 180 self.subbands_num = 4181 self.K = 3 # outputs: |M|, cos∠M, sin∠M182 183 self.downsample_ratio = 2 ** 6 # This number equals 2^{#encoder_blcoks}184 185 self.stft = STFT(186 n_fft=window_size,187 hop_length=hop_size,188 win_length=window_size,189 window=window,190 center=center,191 pad_mode=pad_mode,192 freeze_parameters=True,193 )194 195 self.istft = ISTFT(196 n_fft=window_size,197 hop_length=hop_size,198 win_length=window_size,199 window=window,200 center=center,201 pad_mode=pad_mode,202 freeze_parameters=True,203 )204 205 self.bn0 = nn.BatchNorm2d(window_size // 2 + 1, momentum=momentum)206 207 self.subband = Subband(subbands_num=self.subbands_num)208 209 self.encoder_block1 = EncoderBlock(210 in_channels=input_channels * self.subbands_num,211 out_channels=32,212 condition_size=condition_size,213 kernel_size=(3, 3),214 downsample=(2, 2),215 activation=activation,216 momentum=momentum,217 )218 self.encoder_block2 = EncoderBlock(219 in_channels=32,220 out_channels=64,221 condition_size=condition_size,222 kernel_size=(3, 3),223 downsample=(2, 2),224 activation=activation,225 momentum=momentum,226 )227 self.encoder_block3 = EncoderBlock(228 in_channels=64,229 out_channels=128,230 condition_size=condition_size,231 kernel_size=(3, 3),232 downsample=(2, 2),233 activation=activation,234 momentum=momentum,235 )236 self.encoder_block4 = EncoderBlock(237 in_channels=128,238 out_channels=256,239 condition_size=condition_size,240 kernel_size=(3, 3),241 downsample=(2, 2),242 activation=activation,243 momentum=momentum,244 )245 self.encoder_block5 = EncoderBlock(246 in_channels=256,247 out_channels=384,248 condition_size=condition_size,249 kernel_size=(3, 3),250 downsample=(2, 2),251 activation=activation,252 momentum=momentum,253 )254 self.encoder_block6 = EncoderBlock(255 in_channels=384,256 out_channels=384,257 condition_size=condition_size,258 kernel_size=(3, 3),259 downsample=(2, 2),260 activation=activation,261 momentum=momentum,262 )263 self.conv_block7 = ConvBlock(264 in_channels=384,265 out_channels=384,266 condition_size=condition_size,267 kernel_size=(3, 3),268 activation=activation,269 momentum=momentum,270 )271 self.decoder_block1 = DecoderBlock(272 in_channels=384,273 out_channels=384,274 condition_size=condition_size,275 kernel_size=(3, 3),276 upsample=(2, 2),277 activation=activation,278 momentum=momentum,279 )280 self.decoder_block2 = DecoderBlock(281 in_channels=384,282 out_channels=384,283 condition_size=condition_size,284 kernel_size=(3, 3),285 upsample=(2, 2),286 activation=activation,287 momentum=momentum,288 )289 self.decoder_block3 = DecoderBlock(290 in_channels=384,291 out_channels=256,292 condition_size=condition_size,293 kernel_size=(3, 3),294 upsample=(2, 2),295 activation=activation,296 momentum=momentum,297 )298 self.decoder_block4 = DecoderBlock(299 in_channels=256,300 out_channels=128,301 condition_size=condition_size,302 kernel_size=(3, 3),303 upsample=(2, 2),304 activation=activation,305 momentum=momentum,306 )307 self.decoder_block5 = DecoderBlock(308 in_channels=128,309 out_channels=64,310 condition_size=condition_size,311 kernel_size=(3, 3),312 upsample=(2, 2),313 activation=activation,314 momentum=momentum,315 )316 self.decoder_block6 = DecoderBlock(317 in_channels=64,318 out_channels=32,319 condition_size=condition_size,320 kernel_size=(3, 3),321 upsample=(2, 2),322 activation=activation,323 momentum=momentum,324 )325 326 self.after_conv_block1 = ConvBlock(327 in_channels=32,328 out_channels=32,329 condition_size=condition_size,330 kernel_size=(3, 3),331 activation=activation,332 momentum=momentum,333 )334 335 self.after_conv2 = nn.Conv2d(336 in_channels=32,337 out_channels=input_channels338 * self.subbands_num339 * self.output_sources_num340 * self.K,341 kernel_size=(1, 1),342 stride=(1, 1),343 padding=(0, 0),344 bias=True,345 )346 347 self.init_weights()348 349 def init_weights(self):350 init_bn(self.bn0)351 init_layer(self.after_conv2)352 353 def feature_maps_to_wav(self, x, sp, sin_in, cos_in, audio_length):354 355 batch_size, _, time_steps, freq_bins = x.shape356 357 x = x.reshape(358 batch_size,359 self.output_sources_num,360 self.input_channels,361 self.K,362 time_steps,363 freq_bins,364 )365 # x: (batch_size, output_sources_num, input_channles, K, time_steps, freq_bins)366 367 mask_mag = torch.sigmoid(x[:, :, :, 0, :, :])368 _mask_real = torch.tanh(x[:, :, :, 1, :, :])369 _mask_imag = torch.tanh(x[:, :, :, 2, :, :])370 _, mask_cos, mask_sin = magphase(_mask_real, _mask_imag)371 # mask_cos, mask_sin: (batch_size, output_sources_num, input_channles, time_steps, freq_bins)372 373 # Y = |Y|cos∠Y + j|Y|sin∠Y374 # = |Y|cos(∠X + ∠M) + j|Y|sin(∠X + ∠M)375 # = |Y|(cos∠X cos∠M - sin∠X sin∠M) + j|Y|(sin∠X cos∠M + cos∠X sin∠M)376 out_cos = (377 cos_in[:, None, :, :, :] * mask_cos - sin_in[:, None, :, :, :] * mask_sin378 )379 out_sin = (380 sin_in[:, None, :, :, :] * mask_cos + cos_in[:, None, :, :, :] * mask_sin381 )382 # out_cos, out_sin: (batch_size, output_sources_num, input_channles, time_steps, freq_bins)383 384 # Calculate |Y|.385 out_mag = F.relu_(sp[:, None, :, :, :] * mask_mag)386 # out_mag: (batch_size, output_sources_num, input_channles, time_steps, freq_bins)387 388 # Calculate Y_{real} and Y_{imag} for ISTFT.389 out_real = out_mag * out_cos390 out_imag = out_mag * out_sin391 # out_real, out_imag: (batch_size, output_sources_num, input_channles, time_steps, freq_bins)392 393 # Reformat shape to (n, 1, time_steps, freq_bins) for ISTFT.394 shape = (395 batch_size * self.output_sources_num * self.input_channels,396 1,397 time_steps,398 freq_bins,399 )400 out_real = out_real.reshape(shape)401 out_imag = out_imag.reshape(shape)402 403 # ISTFT.404 wav_out = self.istft(out_real, out_imag, audio_length)405 # (batch_size * output_sources_num * input_channels, segments_num)406 407 # Reshape.408 wav_out = wav_out.reshape(409 batch_size, self.output_sources_num * self.input_channels, audio_length410 )411 # (batch_size, output_sources_num * input_channels, segments_num)412 413 return wav_out414 415 def forward(self, input_dict):416 """417 Args:418 input: (batch_size, segment_samples, channels_num)419 420 Outputs:421 output_dict: {422 'wav': (batch_size, segment_samples, channels_num),423 'sp': (batch_size, channels_num, time_steps, freq_bins)}424 """425 426 mixture = input_dict['waveform']427 condition = input_dict['condition']428 429 sp, cos_in, sin_in = self.wav_to_spectrogram_phase(mixture)430 """(batch_size, channels_num, time_steps, freq_bins)"""431 432 # Batch normalization433 x = sp.transpose(1, 3)434 x = self.bn0(x)435 x = x.transpose(1, 3)436 """(batch_size, chanenls, time_steps, freq_bins)"""437 438 # Pad spectrogram to be evenly divided by downsample ratio.439 origin_len = x.shape[2]440 pad_len = (441 int(np.ceil(x.shape[2] / self.downsample_ratio)) * self.downsample_ratio442 - origin_len443 )444 x = F.pad(x, pad=(0, 0, 0, pad_len))445 """(batch_size, channels, padded_time_steps, freq_bins)"""446 447 # Let frequency bins be evenly divided by 2, e.g., 513 -> 512448 x = x[..., 0 : x.shape[-1] - 1] # (bs, channels, T, F)449 450 x = self.subband.analysis(x)451 452 # UNet453 (x1_pool, x1) = self.encoder_block1(454 x, condition455 ) # x1_pool: (bs, 32, T / 2, F / 2)456 (x2_pool, x2) = self.encoder_block2(457 x1_pool, condition458 ) # x2_pool: (bs, 64, T / 4, F / 4)459 (x3_pool, x3) = self.encoder_block3(460 x2_pool, condition461 ) # x3_pool: (bs, 128, T / 8, F / 8)462 (x4_pool, x4) = self.encoder_block4(463 x3_pool, condition464 ) # x4_pool: (bs, 256, T / 16, F / 16)465 (x5_pool, x5) = self.encoder_block5(466 x4_pool, condition467 ) # x5_pool: (bs, 512, T / 32, F / 32)468 (x6_pool, x6) = self.encoder_block6(469 x5_pool, condition470 ) # x6_pool: (bs, 1024, T / 64, F / 64)471 x_center = self.conv_block7(x6_pool, condition) # (bs, 2048, T / 64, F / 64)472 x7 = self.decoder_block1(x_center, x6, condition) # (bs, 1024, T / 32, F / 32)473 x8 = self.decoder_block2(x7, x5, condition) # (bs, 512, T / 16, F / 16)474 x9 = self.decoder_block3(x8, x4, condition) # (bs, 256, T / 8, F / 8)475 x10 = self.decoder_block4(x9, x3, condition) # (bs, 128, T / 4, F / 4)476 x11 = self.decoder_block5(x10, x2, condition) # (bs, 64, T / 2, F / 2)477 x12 = self.decoder_block6(x11, x1, condition) # (bs, 32, T, F)478 x = self.after_conv_block1(x12, condition) # (bs, 32, T, F)479 x = self.after_conv2(x)480 # (batch_size, input_channles * subbands_num * targets_num * k, T, F // subbands_num)481 482 x = self.subband.synthesis(x)483 # (batch_size, input_channles * targets_num * K, T, F)484 485 # Recover shape486 x = F.pad(x, pad=(0, 1)) # Pad frequency, e.g., 1024 -> 1025.487 x = x[:, :, 0:origin_len, :] # (bs, feature_maps, T, F)488 489 audio_length = mixture.shape[2]490 491 separated_audio = self.feature_maps_to_wav(x, sp, sin_in, cos_in, audio_length)492 # separated_audio: (batch_size, output_sources_num * input_channels, segments_num)493 494 output_dict = {'waveform': separated_audio}495 496 return output_dict497