Spaces:
Running
Running
| import torch | |
| from deepafx_st.models.mobilenetv2 import MobileNetV2 | |
| from deepafx_st.models.efficient_net import EfficientNet | |
| class SpectralEncoder(torch.nn.Module): | |
| def __init__( | |
| self, | |
| num_params, | |
| sample_rate, | |
| encoder_model="mobilenet_v2", | |
| embed_dim=1028, | |
| width_mult=1, | |
| min_level_db=-80, | |
| ): | |
| """Encoder operating on spectrograms. | |
| Args: | |
| num_params (int): Number of processor parameters to generate. | |
| sample_rate (float): Audio sample rate for computing melspectrogram. | |
| encoder_model (str, optional): Encoder model architecture. Default: "mobilenet_v2" | |
| embed_dim (int, optional): Dimentionality of the encoder representations. | |
| width_mult (int, optional): Encoder size. Default: 1 | |
| min_level_db (float, optional): Minimal dB value for the spectrogram. Default: -80 | |
| """ | |
| super().__init__() | |
| self.num_params = num_params | |
| self.sample_rate = sample_rate | |
| self.encoder_model = encoder_model | |
| self.embed_dim = embed_dim | |
| self.width_mult = width_mult | |
| self.min_level_db = min_level_db | |
| # load model from torch.hub | |
| if encoder_model == "mobilenet_v2": | |
| self.encoder = MobileNetV2(embed_dim=embed_dim, width_mult=width_mult) | |
| elif encoder_model == "efficient_net": | |
| self.encoder = EfficientNet.from_name( | |
| "efficientnet-b2", | |
| in_channels=1, | |
| image_size=(128, 65), | |
| include_top=False, | |
| ) | |
| self.embedding_projection = torch.nn.Conv2d( | |
| in_channels=1408, | |
| out_channels=embed_dim, | |
| kernel_size=(1, 1), | |
| stride=(1, 1), | |
| padding=(0, 0), | |
| bias=True, | |
| ) | |
| else: | |
| raise ValueError(f"Invalid encoder_model: {encoder_model}.") | |
| self.window = torch.nn.Parameter(torch.hann_window(4096)) | |
| def forward(self, x): | |
| """ | |
| Args: | |
| x (Tensor): Input waveform of shape [batch x channels x samples] | |
| Returns: | |
| e (Tensor): Latent embedding produced by Encoder. [batch x embed_dim] | |
| """ | |
| bs, chs, samp = x.size() | |
| # compute spectrogram of waveform | |
| X = torch.stft( | |
| x.view(bs, -1), | |
| 4096, | |
| 2048, | |
| window=self.window, | |
| return_complex=True, | |
| ) | |
| X_db = torch.pow(X.abs() + 1e-8, 0.3) | |
| X_db_norm = X_db | |
| # standardize (0, 1) 0.322970 0.278452 | |
| X_db_norm -= 0.322970 | |
| X_db_norm /= 0.278452 | |
| X_db_norm = X_db_norm.unsqueeze(1).permute(0, 1, 3, 2) | |
| if self.encoder_model == "mobilenet_v2": | |
| # repeat channels by 3 to fit vision model | |
| X_db_norm = X_db_norm.repeat(1, 3, 1, 1) | |
| # pass melspectrogram through encoder | |
| e = self.encoder(X_db_norm) | |
| # apply avg pooling across time for encoder embeddings | |
| e = torch.nn.functional.adaptive_avg_pool2d(e, 1).reshape(e.shape[0], -1) | |
| # normalize by L2 norm | |
| norm = torch.norm(e, p=2, dim=-1, keepdim=True) | |
| e_norm = e / norm | |
| elif self.encoder_model == "efficient_net": | |
| # Efficient Net internal downsamples by 32 on time and freq axis, then average pools the rest | |
| e = self.encoder(X_db_norm) | |
| # Adding 1x1 conv to project down or up to the requested embedding size | |
| e = self.embedding_projection(e) | |
| e = torch.squeeze(e, dim=3) | |
| e = torch.squeeze(e, dim=2) | |
| # normalize by L2 norm | |
| norm = torch.norm(e, p=2, dim=-1, keepdim=True) | |
| e_norm = e / norm | |
| return e_norm | |