| 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 |
| 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] |
| |
| """ |
| |
| 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 |
|
|
