Spaces:
Running
Running
| import numpy as np | |
| import torch.nn.functional as F | |
| from torch import nn | |
| from .model import MLPLayers | |
| class LinearProbe(nn.Module): | |
| def __init__(self, model, mlp, freeze, in_ch, out_ch, act=None): | |
| """ | |
| Args: | |
| model: nn.Module | |
| mlp: bool, if True, then use the MLP layer as the linear probe module | |
| freeze: bool, if Ture, then freeze all the CLAP model's layers when training the linear probe | |
| in_ch: int, the output channel from CLAP model | |
| out_ch: int, the output channel from linear probe (class_num) | |
| act: torch.nn.functional, the activation function before the loss function | |
| """ | |
| super().__init__() | |
| in_ch = 512 | |
| self.clap_model = model | |
| self.clap_model.text_branch = None # to save memory | |
| self.freeze = freeze | |
| if mlp: | |
| self.lp_layer = MLPLayers(units=[in_ch, in_ch * 2, out_ch]) | |
| else: | |
| self.lp_layer = nn.Linear(in_ch, out_ch) | |
| if self.freeze: | |
| for param in self.clap_model.parameters(): | |
| param.requires_grad = False | |
| if act == "None": | |
| self.act = None | |
| elif act == "relu": | |
| self.act = nn.ReLU() | |
| elif act == "elu": | |
| self.act = nn.ELU() | |
| elif act == "prelu": | |
| self.act = nn.PReLU(num_parameters=in_ch) | |
| elif act == "softmax": | |
| self.act = nn.Softmax(dim=-1) | |
| elif act == "sigmoid": | |
| self.act = nn.Sigmoid() | |
| def forward(self, x, mix_lambda=None, device=None): | |
| """ | |
| Args: | |
| x: waveform, torch.tensor [batch, t_samples] / batch of mel_spec and longer list | |
| mix_lambda: torch.tensor [batch], the mixup lambda | |
| Returns: | |
| class_prob: torch.tensor [batch, class_num] | |
| """ | |
| # batchnorm cancel grandient | |
| if self.freeze: | |
| self.clap_model.eval() | |
| x = self.clap_model.audio_projection( | |
| self.clap_model.audio_branch(x, mixup_lambda=mix_lambda, device=device)[ | |
| "embedding" | |
| ] | |
| ) | |
| out = self.lp_layer(x) | |
| if self.act is not None: | |
| out = self.act(out) | |
| return out | |