""" /* Copyright (c) 2023 Amazon Written by Jan Buethe */ /* Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: - Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. - Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */ """ import math as m import copy import torch import torch.nn.functional as F from torch import nn from torch.nn.utils import weight_norm, spectral_norm import torchaudio from utils.spec import gen_filterbank # auxiliary functions def remove_all_weight_norms(module): for m in module.modules(): if hasattr(m, 'weight_v'): nn.utils.remove_weight_norm(m) def create_smoothing_kernel(h, w, gamma=1.5): ch = h / 2 - 0.5 cw = w / 2 - 0.5 sh = gamma * ch sw = gamma * cw vx = ((torch.arange(h) - ch) / sh) ** 2 vy = ((torch.arange(w) - cw) / sw) ** 2 vals = vx.view(-1, 1) + vy.view(1, -1) kernel = torch.exp(- vals) kernel = kernel / kernel.sum() return kernel def create_kernel(h, w, sh, sw): # proto kernel gives disjoint partition of 1 proto_kernel = torch.ones((sh, sw)) # create smoothing kernel eta h_eta, w_eta = h - sh + 1, w - sw + 1 assert h_eta > 0 and w_eta > 0 eta = create_smoothing_kernel(h_eta, w_eta).view(1, 1, h_eta, w_eta) kernel0 = F.pad(proto_kernel, [w_eta - 1, w_eta - 1, h_eta - 1, h_eta - 1]).unsqueeze(0).unsqueeze(0) kernel = F.conv2d(kernel0, eta) return kernel # positional embeddings class FrequencyPositionalEmbedding(nn.Module): def __init__(self): super().__init__() def forward(self, x): N = x.size(2) args = torch.arange(0, N, dtype=x.dtype, device=x.device) * torch.pi * 2 / N cos = torch.cos(args).reshape(1, 1, -1, 1) sin = torch.sin(args).reshape(1, 1, -1, 1) zeros = torch.zeros_like(x[:, 0:1, :, :]) y = torch.cat((x, zeros + sin, zeros + cos), dim=1) return y class PositionalEmbedding2D(nn.Module): def __init__(self, d=5): super().__init__() self.d = d def forward(self, x): N = x.size(2) M = x.size(3) h_args = torch.arange(0, N, dtype=x.dtype, device=x.device).reshape(1, 1, -1, 1) w_args = torch.arange(0, M, dtype=x.dtype, device=x.device).reshape(1, 1, 1, -1) coeffs = (10000 ** (-2 * torch.arange(0, self.d, dtype=x.dtype, device=x.device) / self.d)).reshape(1, -1, 1, 1) h_sin = torch.sin(coeffs * h_args) h_cos = torch.sin(coeffs * h_args) w_sin = torch.sin(coeffs * w_args) w_cos = torch.sin(coeffs * w_args) zeros = torch.zeros_like(x[:, 0:1, :, :]) y = torch.cat((x, zeros + h_sin, zeros + h_cos, zeros + w_sin, zeros + w_cos), dim=1) return y # spectral discriminator base class class SpecDiscriminatorBase(nn.Module): RECEPTIVE_FIELD_MAX_WIDTH=10000 def __init__(self, layers, resolution, fs=16000, freq_roi=[50, 7000], noise_gain=1e-3, fmap_start_index=0 ): super().__init__() self.layers = nn.ModuleList(layers) self.resolution = resolution self.fs = fs self.noise_gain = noise_gain self.fmap_start_index = fmap_start_index if fmap_start_index >= len(layers): raise ValueError(f'fmap_start_index is larger than number of layers') # filter bank for noise shaping n_fft = resolution[0] self.filterbank = nn.Parameter( gen_filterbank(n_fft // 2, fs, keep_size=True), requires_grad=False ) # roi bins f_step = fs / n_fft self.start_bin = int(m.ceil(freq_roi[0] / f_step - 0.01)) self.stop_bin = min(int(m.floor(freq_roi[1] / f_step + 0.01)), n_fft//2 + 1) self.init_weights() # determine receptive field size, offsets and strides hw = 1000 while True: x = torch.zeros((1, hw, hw)) with torch.no_grad(): y = self.run_layer_stack(x)[-1] pos0 = [y.size(-2) // 2, y.size(-1) // 2] pos1 = [t + 1 for t in pos0] hs0, ws0 = self._receptive_field((hw, hw), pos0) hs1, ws1 = self._receptive_field((hw, hw), pos1) h0 = hs0[1] - hs0[0] + 1 h1 = hs1[1] - hs1[0] + 1 w0 = ws0[1] - ws0[0] + 1 w1 = ws1[1] - ws1[0] + 1 if h0 != h1 or w0 != w1: hw = 2 * hw else: # strides sh = hs1[0] - hs0[0] sw = ws1[0] - ws0[0] if sh == 0 or sw == 0: continue # offsets oh = hs0[0] - sh * pos0[0] ow = ws0[0] - sw * pos0[1] # overlap factor overlap = w0 / sw + h0 / sh #print(f"{w0=} {h0=} {sw=} {sh=} {overlap=}") self.receptive_field_params = {'width': [sw, ow, w0], 'height': [sh, oh, h0], 'overlap': overlap} break if hw > self.RECEPTIVE_FIELD_MAX_WIDTH: print("warning: exceeded max size while trying to determine receptive field") # create transposed convolutional kernel #self.tconv_kernel = nn.Parameter(create_kernel(h0, w0, sw, sw), requires_grad=False) def run_layer_stack(self, spec): output = [] x = spec.unsqueeze(1) for layer in self.layers: x = layer(x) output.append(x) return output def forward(self, x): """ returns array with feature maps and final score at index -1 """ output = [] x = self.spectrogram(x) output = self.run_layer_stack(x) return output[self.fmap_start_index:] def receptive_field(self, output_pos): if self.receptive_field_params is not None: s, o, h = self.receptive_field_params['height'] h_min = output_pos[0] * s + o + self.start_bin h_max = h_min + h h_min = max(h_min, self.start_bin) h_max = min(h_max, self.stop_bin) s, o, w = self.receptive_field_params['width'] w_min = output_pos[1] * s + o w_max = w_min + w return (h_min, h_max), (w_min, w_max) else: return None, None def _receptive_field(self, input_dims, output_pos): """ determines receptive field probabilistically via autograd (slow) """ x = torch.randn((1,) + input_dims, requires_grad=True) # run input through layers y = self.run_layer_stack(x)[-1] b, c, h, w = y.shape if output_pos[0] >= h or output_pos[1] >= w: raise ValueError("position out of range") mask = torch.zeros((b, c, h, w)) mask[0, 0, output_pos[0], output_pos[1]] = 1 (mask * y).sum().backward() hs, ws = torch.nonzero(x.grad[0], as_tuple=True) h_min, h_max = hs.min().item(), hs.max().item() w_min, w_max = ws.min().item(), ws.max().item() return [h_min, h_max], [w_min, w_max] def init_weights(self): for m in self.modules(): if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d) or isinstance(m, nn.Linear) or isinstance(m, nn.Embedding): nn.init.orthogonal_(m.weight.data) def spectrogram(self, x): n_fft, hop_length, win_length = self.resolution x = x.squeeze(1) window = getattr(torch, 'hann_window')(win_length).to(x.device) x = torch.stft(x, n_fft=n_fft, hop_length=hop_length, win_length=win_length,\ window=window, return_complex=True) #[B, F, T] x = torch.abs(x) # noise floor following spectral envelope smoothed_x = torch.matmul(self.filterbank, x) noise = torch.randn_like(x) * smoothed_x * self.noise_gain x = x + noise # frequency ROI x = x[:, self.start_bin : self.stop_bin + 1, ...] return torchaudio.functional.amplitude_to_DB(x,db_multiplier=0.0, multiplier=20,amin=1e-05,top_db=80)#torch.sqrt(x) def grad_map(self, x): self.zero_grad() n_fft, hop_length, win_length = self.resolution window = getattr(torch, 'hann_window')(win_length).to(x.device) y = torch.stft(x.squeeze(1), n_fft=n_fft, hop_length=hop_length, win_length=win_length, window=window, return_complex=True) #[B, F, T] y = torch.abs(y) specgram = torchaudio.functional.amplitude_to_DB(y,db_multiplier=0.0, multiplier=20,amin=1e-05,top_db=80) specgram.requires_grad = True specgram.retain_grad() if specgram.grad is not None: specgram.grad.zero_() y = specgram[:, self.start_bin : self.stop_bin + 1, ...] scores = self.run_layer_stack(y)[-1] loss = torch.mean((1 - scores) ** 2) loss.backward() return specgram.data[0], torch.abs(specgram.grad)[0] def relevance_map(self, x): n_fft, hop_length, win_length = self.resolution y = x.view(-1) window = getattr(torch, 'hann_window')(win_length).to(x.device) y = torch.stft(y, n_fft=n_fft, hop_length=hop_length, win_length=win_length,\ window=window, return_complex=True) #[B, F, T] y = torch.abs(y) specgram = torchaudio.functional.amplitude_to_DB(y,db_multiplier=0.0, multiplier=20,amin=1e-05,top_db=80) scores = self.forward(x)[-1] sh, _, h = self.receptive_field_params['height'] sw, _, w = self.receptive_field_params['width'] kernel = create_kernel(h, w, sh, sw).float().to(scores.device) with torch.no_grad(): pad_w = (w + sw - 1) // sw pad_h = (h + sh - 1) // sh padded_scores = F.pad(scores, (pad_w, pad_w, pad_h, pad_h), mode='replicate') # CAVE: padding should be derived from offsets rv = F.conv_transpose2d(padded_scores, kernel, bias=None, stride=(sh, sw), padding=(h//2, w//2)) rv = rv[..., pad_h * sh : - pad_h * sh, pad_w * sw : -pad_w * sw] relevance = torch.zeros_like(specgram) relevance[..., self.start_bin : self.start_bin + rv.size(-2), : rv.size(-1)] = rv return specgram, relevance def lrp(self, x, eps=1e-9, label='both', threshold=0.5, low=None, high=None, verbose=False): """ layer-wise relevance propagation (https://git.tu-berlin.de/gmontavon/lrp-tutorial) """ # ToDo: this code is highly unsafe as it assumes that layers are nn.Sequential with suitable activations def newconv2d(layer,g): new_layer = nn.Conv2d(layer.in_channels, layer.out_channels, layer.kernel_size, stride=layer.stride, padding=layer.padding, dilation=layer.dilation, groups=layer.groups) try: new_layer.weight = nn.Parameter(g(layer.weight.data.clone())) except AttributeError: pass try: new_layer.bias = nn.Parameter(g(layer.bias.data.clone())) except AttributeError: pass return new_layer bounds = { 64: [-85.82449722290039, 2.1755014657974243], 128: [-84.49211349487305, 3.5078893899917607], 256: [-80.33127822875977, 7.6687201976776125], 512: [-73.79328079223633, 14.20672025680542], 1024: [-67.59239501953125, 20.40760498046875], 2048: [-62.31902580261231, 25.680974197387698], } nfft = self.resolution[0] if low is None: low = bounds[nfft][0] if high is None: high = bounds[nfft][1] remove_all_weight_norms(self) for p in self.parameters(): if p.grad is not None: p.grad.zero_() num_layers = len(self.layers) X = self.spectrogram(x). detach() # forward pass A = [X.unsqueeze(1)] + [None] * len(self.layers) for i in range(num_layers - 1): A[i + 1] = self.layers[i](A[i]) # initial relevance is last layer without activation r = A[-2] last_layer_rs = [r] layer = self.layers[-1] for sublayer in list(layer)[:-1]: r = sublayer(r) last_layer_rs.append(r) mask = torch.zeros_like(r) mask.requires_grad_(False) if verbose: print(r.min(), r.max()) if label in {'both', 'fake'}: mask[r < -threshold] = 1 if label in {'both', 'real'}: mask[r > threshold] = 1 r = r * mask # backward pass R = [None] * num_layers + [r] for l in range(1, num_layers)[::-1]: A[l] = (A[l]).data.requires_grad_(True) layer = nn.Sequential(*(list(self.layers[l])[:-1])) z = layer(A[l]) + eps s = (R[l+1] / z).data (z*s).sum().backward() c = A[l].grad R[l] = (A[l] * c).data # first layer A[0] = (A[0].data).requires_grad_(True) Xl = (torch.zeros_like(A[0].data) + low).requires_grad_(True) Xh = (torch.zeros_like(A[0].data) + high).requires_grad_(True) if len(list(self.layers)) > 2: # unsafe way to check for embedding layer embed = list(self.layers[0])[0] conv = list(self.layers[0])[1] layer = nn.Sequential(embed, conv) layerl = nn.Sequential(embed, newconv2d(conv, lambda p: p.clamp(min=0))) layerh = nn.Sequential(embed, newconv2d(conv, lambda p: p.clamp(max=0))) else: layer = list(self.layers[0])[0] layerl = newconv2d(layer, lambda p: p.clamp(min=0)) layerh = newconv2d(layer, lambda p: p.clamp(max=0)) z = layer(A[0]) z -= layerl(Xl) + layerh(Xh) s = (R[1] / z).data (z * s).sum().backward() c, cp, cm = A[0].grad, Xl.grad, Xh.grad R[0] = (A[0] * c + Xl * cp + Xh * cm) #R[0] = (A[0] * c).data return X, R[0].mean(dim=1) def create_3x3_conv_plan(num_layers : int, f_stretch : int, f_down : int, t_stretch : int, t_down : int ): """ creates a stride, dilation, padding plan for a 2d conv network Args: num_layers (int): number of layers f_stretch (int): log_2 of stretching factor along frequency axis f_down (int): log_2 of downsampling factor along frequency axis t_stretch (int): log_2 of stretching factor along time axis t_down (int): log_2 of downsampling factor along time axis Returns: list(list(tuple)): list containing entries [(stride_t, stride_f), (dilation_t, dilation_f), (padding_t, padding_f)] """ assert num_layers > 0 and t_stretch >= 0 and t_down >= 0 and f_stretch >= 0 and f_down >= 0 assert f_stretch < num_layers and t_stretch < num_layers def process_dimension(n_layers, stretch, down): stack_layers = n_layers - 1 stride_layers = min(min(down, stretch) , stack_layers) dilation_layers = max(min(stack_layers - stride_layers - 1, stretch - stride_layers), 0) final_stride = 2 ** (max(down - stride_layers, 0)) final_dilation = 1 if stride_layers < stack_layers and stretch - stride_layers - dilation_layers > 0: final_dilation = 2 strides, dilations, paddings = [], [], [] processed_layers = 0 current_dilation = 1 for _ in range(stride_layers): # increase receptive field and downsample via stride = 2 strides.append(2) dilations.append(1) paddings.append(1) processed_layers += 1 if processed_layers < stack_layers: strides.append(1) dilations.append(1) paddings.append(1) processed_layers += 1 for _ in range(dilation_layers): # increase receptive field via dilation = 2 strides.append(1) current_dilation *= 2 dilations.append(current_dilation) paddings.append(current_dilation) processed_layers += 1 while processed_layers < n_layers - 1: # fill up with std layers strides.append(1) dilations.append(current_dilation) paddings.append(current_dilation) processed_layers += 1 # final layer strides.append(final_stride) current_dilation * final_dilation dilations.append(current_dilation) paddings.append(current_dilation) processed_layers += 1 assert processed_layers == n_layers return strides, dilations, paddings t_strides, t_dilations, t_paddings = process_dimension(num_layers, t_stretch, t_down) f_strides, f_dilations, f_paddings = process_dimension(num_layers, f_stretch, f_down) plan = [] for i in range(num_layers): plan.append([ (f_strides[i], t_strides[i]), (f_dilations[i], t_dilations[i]), (f_paddings[i], t_paddings[i]), ]) return plan class DiscriminatorExperimental(SpecDiscriminatorBase): def __init__(self, resolution, fs=16000, freq_roi=[50, 7400], noise_gain=0, num_channels=16, max_channels=512, num_layers=5, use_spectral_norm=False): norm_f = weight_norm if use_spectral_norm == False else spectral_norm self.num_channels = num_channels self.num_channels_max = max_channels self.num_layers = num_layers layers = [] stride = (2, 1) padding= (1, 1) in_channels = 1 + 2 out_channels = self.num_channels for _ in range(self.num_layers): layers.append( nn.Sequential( FrequencyPositionalEmbedding(), norm_f(nn.Conv2d(in_channels, out_channels, (3, 3), stride=stride, padding=padding)), nn.ReLU(inplace=True) ) ) in_channels = out_channels + 2 out_channels = min(2 * out_channels, self.num_channels_max) layers.append( nn.Sequential( FrequencyPositionalEmbedding(), norm_f(nn.Conv2d(in_channels, 1, (3, 3), padding=padding)), nn.Sigmoid() ) ) super().__init__(layers=layers, resolution=resolution, fs=fs, freq_roi=freq_roi, noise_gain=noise_gain) # bias biases bias_val = 0.1 with torch.no_grad(): for name, weight in self.named_parameters(): if 'bias' in name: weight = weight + bias_val configs = { 'f_down': { 'stretch' : { 64 : (0, 0), 128: (1, 0), 256: (2, 0), 512: (3, 0), 1024: (4, 0), 2048: (5, 0) }, 'down' : { 64 : (0, 0), 128: (1, 0), 256: (2, 0), 512: (3, 0), 1024: (4, 0), 2048: (5, 0) } }, 'ft_down': { 'stretch' : { 64 : (0, 4), 128: (1, 3), 256: (2, 2), 512: (3, 1), 1024: (4, 0), 2048: (5, 0) }, 'down' : { 64 : (0, 4), 128: (1, 3), 256: (2, 2), 512: (3, 1), 1024: (4, 0), 2048: (5, 0) } }, 'dilated': { 'stretch' : { 64 : (0, 4), 128: (1, 3), 256: (2, 2), 512: (3, 1), 1024: (4, 0), 2048: (5, 0) }, 'down' : { 64 : (0, 0), 128: (0, 0), 256: (0, 0), 512: (0, 0), 1024: (0, 0), 2048: (0, 0) } }, 'mixed': { 'stretch' : { 64 : (0, 4), 128: (1, 3), 256: (2, 2), 512: (3, 1), 1024: (4, 0), 2048: (5, 0) }, 'down' : { 64 : (0, 0), 128: (1, 0), 256: (2, 0), 512: (3, 0), 1024: (4, 0), 2048: (5, 0) } }, } class DiscriminatorMagFree(SpecDiscriminatorBase): def __init__(self, resolution, fs=16000, freq_roi=[50, 7400], noise_gain=0, num_channels=16, max_channels=256, num_layers=5, use_spectral_norm=False, design=None): if design is None: raise ValueError('error: arch required in DiscriminatorMagFree') norm_f = weight_norm if use_spectral_norm == False else spectral_norm stretch = configs[design]['stretch'][resolution[0]] down = configs[design]['down'][resolution[0]] self.num_channels = num_channels self.num_channels_max = max_channels self.num_layers = num_layers self.stretch = stretch self.down = down layers = [] plan = create_3x3_conv_plan(num_layers + 1, stretch[0], down[0], stretch[1], down[1]) in_channels = 1 + 2 out_channels = self.num_channels for i in range(self.num_layers): layers.append( nn.Sequential( FrequencyPositionalEmbedding(), norm_f(nn.Conv2d(in_channels, out_channels, (3, 3), stride=plan[i][0], dilation=plan[i][1], padding=plan[i][2])), nn.ReLU(inplace=True) ) ) in_channels = out_channels + 2 # product over strides channel_factor = plan[i][0][0] * plan[i][0][1] out_channels = min(channel_factor * out_channels, self.num_channels_max) layers.append( nn.Sequential( FrequencyPositionalEmbedding(), norm_f(nn.Conv2d(in_channels, 1, (3, 3), stride=plan[-1][0], dilation=plan[-1][1], padding=plan[-1][2])), nn.Sigmoid() ) ) # for layer in layers: # print(layer) # print("end\n\n") super().__init__(layers=layers, resolution=resolution, fs=fs, freq_roi=freq_roi, noise_gain=noise_gain) # bias biases bias_val = 0.1 with torch.no_grad(): for name, weight in self.named_parameters(): if 'bias' in name: weight = weight + bias_val class DiscriminatorMagFreqPosition(SpecDiscriminatorBase): def __init__(self, resolution, fs=16000, freq_roi=[50, 7400], noise_gain=0, num_channels=16, max_channels=512, num_layers=5, use_spectral_norm=False): norm_f = weight_norm if use_spectral_norm == False else spectral_norm self.num_channels = num_channels self.num_channels_max = max_channels self.num_layers = num_layers layers = [] stride = (2, 1) padding= (1, 1) in_channels = 1 + 2 out_channels = self.num_channels for _ in range(self.num_layers): layers.append( nn.Sequential( FrequencyPositionalEmbedding(), norm_f(nn.Conv2d(in_channels, out_channels, (3, 3), stride=stride, padding=padding)), nn.LeakyReLU(0.2, inplace=True) ) ) in_channels = out_channels + 2 out_channels = min(2 * out_channels, self.num_channels_max) layers.append( nn.Sequential( FrequencyPositionalEmbedding(), norm_f(nn.Conv2d(in_channels, 1, (3, 3), padding=padding)) ) ) super().__init__(layers=layers, resolution=resolution, fs=fs, freq_roi=freq_roi, noise_gain=noise_gain) class DiscriminatorMag2dPositional(SpecDiscriminatorBase): def __init__(self, resolution, fs=16000, freq_roi=[50, 7400], noise_gain=0, num_channels=16, max_channels=512, num_layers=5, d=5, use_spectral_norm=False): norm_f = weight_norm if use_spectral_norm == False else spectral_norm self.resolution = resolution self.num_channels = num_channels self.num_channels_max = max_channels self.num_layers = num_layers self.d = d embedding_dim = 4 * d layers = [] stride = (2, 2) padding= (1, 1) in_channels = 1 + embedding_dim out_channels = self.num_channels for _ in range(self.num_layers): layers.append( nn.Sequential( PositionalEmbedding2D(d), norm_f(nn.Conv2d(in_channels, out_channels, (3, 3), stride=stride, padding=padding)), nn.LeakyReLU(0.2, inplace=True) ) ) in_channels = out_channels + embedding_dim out_channels = min(2 * out_channels, self.num_channels_max) layers.append( nn.Sequential( PositionalEmbedding2D(), norm_f(nn.Conv2d(in_channels, 1, (3, 3), padding=padding)) ) ) super().__init__(layers=layers, resolution=resolution, fs=fs, freq_roi=freq_roi, noise_gain=noise_gain) class DiscriminatorMag(SpecDiscriminatorBase): def __init__(self, resolution, fs=16000, freq_roi=[50, 7400], noise_gain=0, num_channels=32, num_layers=5, use_spectral_norm=False): norm_f = weight_norm if use_spectral_norm == False else spectral_norm self.num_channels = num_channels self.num_layers = num_layers layers = [] stride = (1, 1) padding= (1, 1) in_channels = 1 out_channels = self.num_channels for _ in range(self.num_layers): layers.append( nn.Sequential( norm_f(nn.Conv2d(in_channels, out_channels, (3, 3), stride=stride, padding=padding)), nn.LeakyReLU(0.2, inplace=True) ) ) in_channels = out_channels layers.append(norm_f(nn.Conv2d(in_channels, 1, (3, 3), padding=padding))) super().__init__(layers=layers, resolution=resolution, fs=fs, freq_roi=freq_roi, noise_gain=noise_gain) discriminators = { 'mag': DiscriminatorMag, 'freqpos': DiscriminatorMagFreqPosition, '2dpos': DiscriminatorMag2dPositional, 'experimental': DiscriminatorExperimental, 'free': DiscriminatorMagFree } class TFDMultiResolutionDiscriminator(torch.nn.Module): def __init__(self, fft_sizes_16k=[64, 128, 256, 512, 1024, 2048], architecture='mag', fs=16000, freq_roi=[50, 7400], noise_gain=0, use_spectral_norm=False, **kwargs): super().__init__() fft_sizes = [int(round(fft_size_16k * fs / 16000)) for fft_size_16k in fft_sizes_16k] resolutions = [[n_fft, n_fft // 4, n_fft] for n_fft in fft_sizes] Disc = discriminators[architecture] discs = [Disc(resolutions[i], fs=fs, freq_roi=freq_roi, noise_gain=noise_gain, use_spectral_norm=use_spectral_norm, **kwargs) for i in range(len(resolutions))] self.discriminators = nn.ModuleList(discs) def forward(self, y): outputs = [] for disc in self.discriminators: outputs.append(disc(y)) return outputs class FWGAN_disc_wrapper(nn.Module): def __init__(self, disc): super().__init__() self.disc = disc def forward(self, y, y_hat): out_real = self.disc(y) out_fake = self.disc(y_hat) y_d_rs = [] y_d_gs = [] fmap_rs = [] fmap_gs = [] for y_real, y_fake in zip(out_real, out_fake): y_d_rs.append(y_real[-1]) y_d_gs.append(y_fake[-1]) fmap_rs.append(y_real[:-1]) fmap_gs.append(y_fake[:-1]) return y_d_rs, y_d_gs, fmap_rs, fmap_gs