Spaces:
Sleeping
Sleeping
File size: 2,041 Bytes
9b0d6c2 |
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 |
import torch
from models.beats.BEATs import BEATsConfig, BEATs
from models.transformer_wrapper import BaseModelWrapper
class BEATsWrapper(BaseModelWrapper):
def __init__(self):
super().__init__()
cfg = BEATsConfig()
self.beats = BEATs(cfg)
def mel_forward(self, x):
with torch.autocast(device_type="cuda", enabled=False):
mel = self.beats.preprocess(x)
mel = mel.unsqueeze(1).transpose(2, 3)
return mel
def forward(self, x):
x = x.transpose(2, 3)
features = self.beats.extract_features(x, do_preprocess=False)[0]
return features
def separate_params(self):
pt_params = [[], [], [], [], [], [], [], [], [], [], [], []]
for k, p in self.named_parameters():
if ".layers.0." in k:
pt_params[0].append(p)
elif ".layers.1." in k:
pt_params[1].append(p)
elif ".layers.2." in k:
pt_params[2].append(p)
elif ".layers.3." in k:
pt_params[3].append(p)
elif ".layers.4." in k:
pt_params[4].append(p)
elif ".layers.5." in k:
pt_params[5].append(p)
elif ".layers.6." in k:
pt_params[6].append(p)
elif ".layers.7." in k:
pt_params[7].append(p)
elif ".layers.8." in k:
pt_params[8].append(p)
elif ".layers.9." in k:
pt_params[9].append(p)
elif ".layers.10." in k:
pt_params[10].append(p)
elif ".layers.11." in k:
pt_params[11].append(p)
elif (".post_extract_proj." in k or ".patch_embedding." in k or '.pos_conv.' in k
or 'beats.layer_norm.' in k or "beats.encoder.layer_norm." in k):
pt_params[0].append(p)
else:
raise ValueError(f"Check separate params for BEATs! Unknown key: {k}")
return list(reversed(pt_params))
|