from functools import partial import numpy as np import torch from torch import nn, Tensor from torch.nn.modules.batchnorm import _BatchNorm from .shared import BackboneRegistry, ComplexConv2d, ComplexConvTranspose2d, ComplexLinear, \ DiffusionStepEmbedding, GaussianFourierProjection, FeatureMapDense, torch_complex_from_reim def get_activation(name): if name == "silu": return nn.SiLU elif name == "relu": return nn.ReLU elif name == "leaky_relu": return nn.LeakyReLU else: raise NotImplementedError(f"Unknown activation: {name}") class BatchNorm(_BatchNorm): def _check_input_dim(self, input): if input.dim() < 2 or input.dim() > 4: raise ValueError("expected 4D or 3D input (got {}D input)".format(input.dim())) class OnReIm(nn.Module): def __init__(self, module_cls, *args, **kwargs): super().__init__() self.re_module = module_cls(*args, **kwargs) self.im_module = module_cls(*args, **kwargs) def forward(self, x): return torch_complex_from_reim(self.re_module(x.real), self.im_module(x.imag)) # Code for DCUNet largely copied from Danilo's `informedenh` repo, cheers! def unet_decoder_args(encoders, *, skip_connections): """Get list of decoder arguments for upsampling (right) side of a symmetric u-net, given the arguments used to construct the encoder. Args: encoders (tuple of length `N` of tuples of (in_chan, out_chan, kernel_size, stride, padding)): List of arguments used to construct the encoders skip_connections (bool): Whether to include skip connections in the calculation of decoder input channels. Return: tuple of length `N` of tuples of (in_chan, out_chan, kernel_size, stride, padding): Arguments to be used to construct decoders """ decoder_args = [] for enc_in_chan, enc_out_chan, enc_kernel_size, enc_stride, enc_padding, enc_dilation in reversed(encoders): if skip_connections and decoder_args: skip_in_chan = enc_out_chan else: skip_in_chan = 0 decoder_args.append( (enc_out_chan + skip_in_chan, enc_in_chan, enc_kernel_size, enc_stride, enc_padding, enc_dilation) ) return tuple(decoder_args) def make_unet_encoder_decoder_args(encoder_args, decoder_args): encoder_args = tuple( ( in_chan, out_chan, tuple(kernel_size), tuple(stride), tuple([n // 2 for n in kernel_size]) if padding == "auto" else tuple(padding), tuple(dilation) ) for in_chan, out_chan, kernel_size, stride, padding, dilation in encoder_args ) if decoder_args == "auto": decoder_args = unet_decoder_args( encoder_args, skip_connections=True, ) else: decoder_args = tuple( ( in_chan, out_chan, tuple(kernel_size), tuple(stride), tuple([n // 2 for n in kernel_size]) if padding == "auto" else padding, tuple(dilation), output_padding, ) for in_chan, out_chan, kernel_size, stride, padding, dilation, output_padding in decoder_args ) return encoder_args, decoder_args DCUNET_ARCHITECTURES = { "DCUNet-10": make_unet_encoder_decoder_args( # Encoders: # (in_chan, out_chan, kernel_size, stride, padding, dilation) ( (1, 32, (7, 5), (2, 2), "auto", (1,1)), (32, 64, (7, 5), (2, 2), "auto", (1,1)), (64, 64, (5, 3), (2, 2), "auto", (1,1)), (64, 64, (5, 3), (2, 2), "auto", (1,1)), (64, 64, (5, 3), (2, 1), "auto", (1,1)), ), # Decoders: automatic inverse "auto", ), "DCUNet-16": make_unet_encoder_decoder_args( # Encoders: # (in_chan, out_chan, kernel_size, stride, padding, dilation) ( (1, 32, (7, 5), (2, 2), "auto", (1,1)), (32, 32, (7, 5), (2, 1), "auto", (1,1)), (32, 64, (7, 5), (2, 2), "auto", (1,1)), (64, 64, (5, 3), (2, 1), "auto", (1,1)), (64, 64, (5, 3), (2, 2), "auto", (1,1)), (64, 64, (5, 3), (2, 1), "auto", (1,1)), (64, 64, (5, 3), (2, 2), "auto", (1,1)), (64, 64, (5, 3), (2, 1), "auto", (1,1)), ), # Decoders: automatic inverse "auto", ), "DCUNet-20": make_unet_encoder_decoder_args( # Encoders: # (in_chan, out_chan, kernel_size, stride, padding, dilation) ( (1, 32, (7, 1), (1, 1), "auto", (1,1)), (32, 32, (1, 7), (1, 1), "auto", (1,1)), (32, 64, (7, 5), (2, 2), "auto", (1,1)), (64, 64, (7, 5), (2, 1), "auto", (1,1)), (64, 64, (5, 3), (2, 2), "auto", (1,1)), (64, 64, (5, 3), (2, 1), "auto", (1,1)), (64, 64, (5, 3), (2, 2), "auto", (1,1)), (64, 64, (5, 3), (2, 1), "auto", (1,1)), (64, 64, (5, 3), (2, 2), "auto", (1,1)), (64, 90, (5, 3), (2, 1), "auto", (1,1)), ), # Decoders: automatic inverse "auto", ), "DilDCUNet-v2": make_unet_encoder_decoder_args( # architecture used in SGMSE / Interspeech paper # Encoders: # (in_chan, out_chan, kernel_size, stride, padding, dilation) ( (1, 32, (4, 4), (1, 1), "auto", (1, 1)), (32, 32, (4, 4), (1, 1), "auto", (1, 1)), (32, 32, (4, 4), (1, 1), "auto", (1, 1)), (32, 64, (4, 4), (2, 1), "auto", (2, 1)), (64, 128, (4, 4), (2, 2), "auto", (4, 1)), (128, 256, (4, 4), (2, 2), "auto", (8, 1)), ), # Decoders: automatic inverse "auto", ), } @BackboneRegistry.register("dcunet") class DCUNet(nn.Module): @staticmethod def add_argparse_args(parser): parser.add_argument("--dcunet-architecture", type=str, default="DilDCUNet-v2", choices=DCUNET_ARCHITECTURES.keys(), help="The concrete DCUNet architecture. 'DilDCUNet-v2' by default.") parser.add_argument("--dcunet-time-embedding", type=str, choices=("gfp", "ds", "none"), default="gfp", help="Timestep embedding style. 'gfp' (Gaussian Fourier Projections) by default.") parser.add_argument("--dcunet-temb-layers-global", type=int, default=1, help="Number of global linear+activation layers for the time embedding. 1 by default.") parser.add_argument("--dcunet-temb-layers-local", type=int, default=1, help="Number of local (per-encoder/per-decoder) linear+activation layers for the time embedding. 1 by default.") parser.add_argument("--dcunet-temb-activation", type=str, default="silu", help="The (complex) activation to use between all (global&local) time embedding layers.") parser.add_argument("--dcunet-time-embedding-complex", action="store_true", help="Use complex-valued timestep embedding. Compatible with 'gfp' and 'ds' embeddings.") parser.add_argument("--dcunet-fix-length", type=str, default="pad", choices=("pad", "trim", "none"), help="DCUNet strategy to 'fix' mismatched input timespan. 'pad' by default.") parser.add_argument("--dcunet-mask-bound", type=str, choices=("tanh", "sigmoid", "none"), default="none", help="DCUNet output bounding strategy. 'none' by default.") parser.add_argument("--dcunet-norm-type", type=str, choices=("bN", "CbN"), default="bN", help="The type of norm to use within each encoder and decoder layer. 'bN' (real/imaginary separate batch norm) by default.") parser.add_argument("--dcunet-activation", type=str, choices=("leaky_relu", "relu", "silu"), default="leaky_relu", help="The activation to use within each encoder and decoder layer. 'leaky_relu' by default.") return parser def __init__( self, dcunet_architecture: str = "DilDCUNet-v2", dcunet_time_embedding: str = "gfp", dcunet_temb_layers_global: int = 2, dcunet_temb_layers_local: int = 1, dcunet_temb_activation: str = "silu", dcunet_time_embedding_complex: bool = False, dcunet_fix_length: str = "pad", dcunet_mask_bound: str = "none", dcunet_norm_type: str = "bN", dcunet_activation: str = "relu", embed_dim: int = 128, **kwargs ): super().__init__() self.architecture = dcunet_architecture self.fix_length_mode = (dcunet_fix_length if dcunet_fix_length != "none" else None) self.norm_type = dcunet_norm_type self.activation = dcunet_activation self.input_channels = 2 # for x_t and y -- note that this is 2 rather than 4, because we directly treat complex channels in this DNN self.time_embedding = (dcunet_time_embedding if dcunet_time_embedding != "none" else None) self.time_embedding_complex = dcunet_time_embedding_complex self.temb_layers_global = dcunet_temb_layers_global self.temb_layers_local = dcunet_temb_layers_local self.temb_activation = dcunet_temb_activation conf_encoders, conf_decoders = DCUNET_ARCHITECTURES[dcunet_architecture] # Replace `input_channels` in encoders config _replaced_input_channels, *rest = conf_encoders[0] encoders = ((self.input_channels, *rest), *conf_encoders[1:]) decoders = conf_decoders self.encoders_stride_product = np.prod( [enc_stride for _, _, _, enc_stride, _, _ in encoders], axis=0 ) # Prepare kwargs for encoder and decoder (to potentially be modified before layer instantiation) encoder_decoder_kwargs = dict( norm_type=self.norm_type, activation=self.activation, temb_layers=self.temb_layers_local, temb_activation=self.temb_activation) # Instantiate (global) time embedding layer embed_ops = [] if self.time_embedding is not None: complex_valued = self.time_embedding_complex if self.time_embedding == "gfp": embed_ops += [GaussianFourierProjection(embed_dim=embed_dim, complex_valued=complex_valued)] encoder_decoder_kwargs["embed_dim"] = embed_dim elif self.time_embedding == "ds": embed_ops += [DiffusionStepEmbedding(embed_dim=embed_dim, complex_valued=complex_valued)] encoder_decoder_kwargs["embed_dim"] = embed_dim if self.time_embedding_complex: assert self.time_embedding in ("gfp", "ds"), "Complex timestep embedding only available for gfp and ds" encoder_decoder_kwargs["complex_time_embedding"] = True for _ in range(self.temb_layers_global): embed_ops += [ ComplexLinear(embed_dim, embed_dim, complex_valued=True), OnReIm(get_activation(dcunet_temb_activation)) ] self.embed = nn.Sequential(*embed_ops) ### Instantiate DCUNet layers ### output_layer = ComplexConvTranspose2d(*decoders[-1]) encoders = [DCUNetComplexEncoderBlock(*args, **encoder_decoder_kwargs) for args in encoders] decoders = [DCUNetComplexDecoderBlock(*args, **encoder_decoder_kwargs) for args in decoders[:-1]] self.mask_bound = (dcunet_mask_bound if dcunet_mask_bound != "none" else None) if self.mask_bound is not None: raise NotImplementedError("sorry, mask bounding not implemented at the moment") # TODO we can't use nn.Sequential since the ComplexConvTranspose2d needs a second `output_size` argument #operations = (output_layer, complex_nn.BoundComplexMask(self.mask_bound)) #output_layer = nn.Sequential(*[x for x in operations if x is not None]) assert len(encoders) == len(decoders) + 1 self.encoders = nn.ModuleList(encoders) self.decoders = nn.ModuleList(decoders) self.output_layer = output_layer or nn.Identity() def forward(self, spec, t) -> Tensor: """ Input shape is expected to be $(batch, nfreqs, time)$, with $nfreqs - 1$ divisible by $f_0 * f_1 * ... * f_N$ where $f_k$ are the frequency strides of the encoders, and $time - 1$ is divisible by $t_0 * t_1 * ... * t_N$ where $t_N$ are the time strides of the encoders. Args: spec (Tensor): complex spectrogram tensor. 1D, 2D or 3D tensor, time last. Returns: Tensor, of shape (batch, time) or (time). """ # TF-rep shape: (batch, self.input_channels, n_fft, frames) # Estimate mask from time-frequency representation. x_in = self.fix_input_dims(spec) x = x_in t_embed = self.embed(t+0j) if self.time_embedding is not None else None enc_outs = [] for idx, enc in enumerate(self.encoders): x = enc(x, t_embed) # UNet skip connection enc_outs.append(x) for (enc_out, dec) in zip(reversed(enc_outs[:-1]), self.decoders): x = dec(x, t_embed, output_size=enc_out.shape) x = torch.cat([x, enc_out], dim=1) output = self.output_layer(x, output_size=x_in.shape) # output shape: (batch, 1, n_fft, frames) output = self.fix_output_dims(output, spec) return output def fix_input_dims(self, x): return _fix_dcu_input_dims( self.fix_length_mode, x, torch.from_numpy(self.encoders_stride_product) ) def fix_output_dims(self, out, x): return _fix_dcu_output_dims(self.fix_length_mode, out, x) def _fix_dcu_input_dims(fix_length_mode, x, encoders_stride_product): """Pad or trim `x` to a length compatible with DCUNet.""" freq_prod = int(encoders_stride_product[0]) time_prod = int(encoders_stride_product[1]) if (x.shape[2] - 1) % freq_prod: raise TypeError( f"Input shape must be [batch, ch, freq + 1, time + 1] with freq divisible by " f"{freq_prod}, got {x.shape} instead" ) time_remainder = (x.shape[3] - 1) % time_prod if time_remainder: if fix_length_mode is None: raise TypeError( f"Input shape must be [batch, ch, freq + 1, time + 1] with time divisible by " f"{time_prod}, got {x.shape} instead. Set the 'fix_length_mode' argument " f"in 'DCUNet' to 'pad' or 'trim' to fix shapes automatically." ) elif fix_length_mode == "pad": pad_shape = [0, time_prod - time_remainder] x = nn.functional.pad(x, pad_shape, mode="constant") elif fix_length_mode == "trim": pad_shape = [0, -time_remainder] x = nn.functional.pad(x, pad_shape, mode="constant") else: raise ValueError(f"Unknown fix_length mode '{fix_length_mode}'") return x def _fix_dcu_output_dims(fix_length_mode, out, x): """Fix shape of `out` to the original shape of `x` by padding/cropping.""" inp_len = x.shape[-1] output_len = out.shape[-1] return nn.functional.pad(out, [0, inp_len - output_len]) def _get_norm(norm_type): if norm_type == "CbN": return ComplexBatchNorm elif norm_type == "bN": return partial(OnReIm, BatchNorm) else: raise NotImplementedError(f"Unknown norm type: {norm_type}") class DCUNetComplexEncoderBlock(nn.Module): def __init__( self, in_chan, out_chan, kernel_size, stride, padding, dilation, norm_type="bN", activation="leaky_relu", embed_dim=None, complex_time_embedding=False, temb_layers=1, temb_activation="silu" ): super().__init__() self.in_chan = in_chan self.out_chan = out_chan self.kernel_size = kernel_size self.stride = stride self.padding = padding self.dilation = dilation self.temb_layers = temb_layers self.temb_activation = temb_activation self.complex_time_embedding = complex_time_embedding self.conv = ComplexConv2d( in_chan, out_chan, kernel_size, stride, padding, bias=norm_type is None, dilation=dilation ) self.norm = _get_norm(norm_type)(out_chan) self.activation = OnReIm(get_activation(activation)) self.embed_dim = embed_dim if self.embed_dim is not None: ops = [] for _ in range(max(0, self.temb_layers - 1)): ops += [ ComplexLinear(self.embed_dim, self.embed_dim, complex_valued=True), OnReIm(get_activation(self.temb_activation)) ] ops += [ FeatureMapDense(self.embed_dim, self.out_chan, complex_valued=True), OnReIm(get_activation(self.temb_activation)) ] self.embed_layer = nn.Sequential(*ops) def forward(self, x, t_embed): y = self.conv(x) if self.embed_dim is not None: y = y + self.embed_layer(t_embed) return self.activation(self.norm(y)) class DCUNetComplexDecoderBlock(nn.Module): def __init__( self, in_chan, out_chan, kernel_size, stride, padding, dilation, output_padding=(0, 0), norm_type="bN", activation="leaky_relu", embed_dim=None, temb_layers=1, temb_activation='swish', complex_time_embedding=False, ): super().__init__() self.in_chan = in_chan self.out_chan = out_chan self.kernel_size = kernel_size self.stride = stride self.padding = padding self.dilation = dilation self.output_padding = output_padding self.complex_time_embedding = complex_time_embedding self.temb_layers = temb_layers self.temb_activation = temb_activation self.deconv = ComplexConvTranspose2d( in_chan, out_chan, kernel_size, stride, padding, output_padding, dilation=dilation, bias=norm_type is None ) self.norm = _get_norm(norm_type)(out_chan) self.activation = OnReIm(get_activation(activation)) self.embed_dim = embed_dim if self.embed_dim is not None: ops = [] for _ in range(max(0, self.temb_layers - 1)): ops += [ ComplexLinear(self.embed_dim, self.embed_dim, complex_valued=True), OnReIm(get_activation(self.temb_activation)) ] ops += [ FeatureMapDense(self.embed_dim, self.out_chan, complex_valued=True), OnReIm(get_activation(self.temb_activation)) ] self.embed_layer = nn.Sequential(*ops) def forward(self, x, t_embed, output_size=None): y = self.deconv(x, output_size=output_size) if self.embed_dim is not None: y = y + self.embed_layer(t_embed) return self.activation(self.norm(y)) # From https://github.com/chanil1218/DCUnet.pytorch/blob/2dcdd30804be47a866fde6435cbb7e2f81585213/models/layers/complexnn.py class ComplexBatchNorm(torch.nn.Module): def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=False): super(ComplexBatchNorm, self).__init__() self.num_features = num_features self.eps = eps self.momentum = momentum self.affine = affine self.track_running_stats = track_running_stats if self.affine: self.Wrr = torch.nn.Parameter(torch.Tensor(num_features)) self.Wri = torch.nn.Parameter(torch.Tensor(num_features)) self.Wii = torch.nn.Parameter(torch.Tensor(num_features)) self.Br = torch.nn.Parameter(torch.Tensor(num_features)) self.Bi = torch.nn.Parameter(torch.Tensor(num_features)) else: self.register_parameter('Wrr', None) self.register_parameter('Wri', None) self.register_parameter('Wii', None) self.register_parameter('Br', None) self.register_parameter('Bi', None) if self.track_running_stats: self.register_buffer('RMr', torch.zeros(num_features)) self.register_buffer('RMi', torch.zeros(num_features)) self.register_buffer('RVrr', torch.ones (num_features)) self.register_buffer('RVri', torch.zeros(num_features)) self.register_buffer('RVii', torch.ones (num_features)) self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long)) else: self.register_parameter('RMr', None) self.register_parameter('RMi', None) self.register_parameter('RVrr', None) self.register_parameter('RVri', None) self.register_parameter('RVii', None) self.register_parameter('num_batches_tracked', None) self.reset_parameters() def reset_running_stats(self): if self.track_running_stats: self.RMr.zero_() self.RMi.zero_() self.RVrr.fill_(1) self.RVri.zero_() self.RVii.fill_(1) self.num_batches_tracked.zero_() def reset_parameters(self): self.reset_running_stats() if self.affine: self.Br.data.zero_() self.Bi.data.zero_() self.Wrr.data.fill_(1) self.Wri.data.uniform_(-.9, +.9) # W will be positive-definite self.Wii.data.fill_(1) def _check_input_dim(self, xr, xi): assert(xr.shape == xi.shape) assert(xr.size(1) == self.num_features) def forward(self, x): xr, xi = x.real, x.imag self._check_input_dim(xr, xi) exponential_average_factor = 0.0 if self.training and self.track_running_stats: self.num_batches_tracked += 1 if self.momentum is None: # use cumulative moving average exponential_average_factor = 1.0 / self.num_batches_tracked.item() else: # use exponential moving average exponential_average_factor = self.momentum # # NOTE: The precise meaning of the "training flag" is: # True: Normalize using batch statistics, update running statistics # if they are being collected. # False: Normalize using running statistics, ignore batch statistics. # training = self.training or not self.track_running_stats redux = [i for i in reversed(range(xr.dim())) if i!=1] vdim = [1] * xr.dim() vdim[1] = xr.size(1) # # Mean M Computation and Centering # # Includes running mean update if training and running. # if training: Mr, Mi = xr, xi for d in redux: Mr = Mr.mean(d, keepdim=True) Mi = Mi.mean(d, keepdim=True) if self.track_running_stats: self.RMr.lerp_(Mr.squeeze(), exponential_average_factor) self.RMi.lerp_(Mi.squeeze(), exponential_average_factor) else: Mr = self.RMr.view(vdim) Mi = self.RMi.view(vdim) xr, xi = xr-Mr, xi-Mi # # Variance Matrix V Computation # # Includes epsilon numerical stabilizer/Tikhonov regularizer. # Includes running variance update if training and running. # if training: Vrr = xr * xr Vri = xr * xi Vii = xi * xi for d in redux: Vrr = Vrr.mean(d, keepdim=True) Vri = Vri.mean(d, keepdim=True) Vii = Vii.mean(d, keepdim=True) if self.track_running_stats: self.RVrr.lerp_(Vrr.squeeze(), exponential_average_factor) self.RVri.lerp_(Vri.squeeze(), exponential_average_factor) self.RVii.lerp_(Vii.squeeze(), exponential_average_factor) else: Vrr = self.RVrr.view(vdim) Vri = self.RVri.view(vdim) Vii = self.RVii.view(vdim) Vrr = Vrr + self.eps Vri = Vri Vii = Vii + self.eps # # Matrix Inverse Square Root U = V^-0.5 # # sqrt of a 2x2 matrix, # - https://en.wikipedia.org/wiki/Square_root_of_a_2_by_2_matrix tau = Vrr + Vii delta = torch.addcmul(Vrr * Vii, Vri, Vri, value=-1) s = delta.sqrt() t = (tau + 2*s).sqrt() # matrix inverse, http://mathworld.wolfram.com/MatrixInverse.html rst = (s * t).reciprocal() Urr = (s + Vii) * rst Uii = (s + Vrr) * rst Uri = ( - Vri) * rst # # Optionally left-multiply U by affine weights W to produce combined # weights Z, left-multiply the inputs by Z, then optionally bias them. # # y = Zx + B # y = WUx + B # y = [Wrr Wri][Urr Uri] [xr] + [Br] # [Wir Wii][Uir Uii] [xi] [Bi] # if self.affine: Wrr, Wri, Wii = self.Wrr.view(vdim), self.Wri.view(vdim), self.Wii.view(vdim) Zrr = (Wrr * Urr) + (Wri * Uri) Zri = (Wrr * Uri) + (Wri * Uii) Zir = (Wri * Urr) + (Wii * Uri) Zii = (Wri * Uri) + (Wii * Uii) else: Zrr, Zri, Zir, Zii = Urr, Uri, Uri, Uii yr = (Zrr * xr) + (Zri * xi) yi = (Zir * xr) + (Zii * xi) if self.affine: yr = yr + self.Br.view(vdim) yi = yi + self.Bi.view(vdim) return torch.view_as_complex(torch.stack([yr, yi], dim=-1)) def extra_repr(self): return '{num_features}, eps={eps}, momentum={momentum}, affine={affine}, ' \ 'track_running_stats={track_running_stats}'.format(**self.__dict__)