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