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