File size: 2,243 Bytes
f1069cc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
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
|