| | import torch |
| | import fairseq |
| | from packaging import version |
| | import torch.nn.functional as F |
| | from fairseq import tasks |
| | from fairseq.checkpoint_utils import load_checkpoint_to_cpu |
| | from fairseq.dataclass.utils import convert_namespace_to_omegaconf |
| | from omegaconf import OmegaConf |
| | from s3prl.upstream.interfaces import UpstreamBase |
| | from torch.nn.utils.rnn import pad_sequence |
| |
|
| | def load_model(filepath): |
| | state = torch.load(filepath, map_location=lambda storage, loc: storage) |
| | |
| | state["cfg"] = OmegaConf.create(state["cfg"]) |
| |
|
| | if "args" in state and state["args"] is not None: |
| | cfg = convert_namespace_to_omegaconf(state["args"]) |
| | elif "cfg" in state and state["cfg"] is not None: |
| | cfg = state["cfg"] |
| | else: |
| | raise RuntimeError( |
| | f"Neither args nor cfg exist in state keys = {state.keys()}" |
| | ) |
| |
|
| | task = tasks.setup_task(cfg.task) |
| | if "task_state" in state: |
| | task.load_state_dict(state["task_state"]) |
| |
|
| | model = task.build_model(cfg.model) |
| |
|
| | return model, cfg, task |
| |
|
| |
|
| | |
| | |
| | |
| | class UpstreamExpert(UpstreamBase): |
| | def __init__(self, ckpt, **kwargs): |
| | super().__init__(**kwargs) |
| | assert version.parse(fairseq.__version__) > version.parse( |
| | "0.10.2" |
| | ), "Please install the fairseq master branch." |
| |
|
| | model, cfg, task = load_model(ckpt) |
| | self.model = model |
| | self.task = task |
| |
|
| | if len(self.hooks) == 0: |
| | module_name = "self.model.encoder.layers" |
| | for module_id in range(len(eval(module_name))): |
| | self.add_hook( |
| | f"{module_name}[{module_id}]", |
| | lambda input, output: input[0].transpose(0, 1), |
| | ) |
| | self.add_hook("self.model.encoder", lambda input, output: output[0]) |
| |
|
| | def forward(self, wavs): |
| | if self.task.cfg.normalize: |
| | wavs = [F.layer_norm(wav, wav.shape) for wav in wavs] |
| |
|
| | device = wavs[0].device |
| | wav_lengths = torch.LongTensor([len(wav) for wav in wavs]).to(device) |
| | wav_padding_mask = ~torch.lt( |
| | torch.arange(max(wav_lengths)).unsqueeze(0).to(device), |
| | wav_lengths.unsqueeze(1), |
| | ) |
| | padded_wav = pad_sequence(wavs, batch_first=True) |
| |
|
| | features, feat_padding_mask = self.model.extract_features( |
| | padded_wav, |
| | padding_mask=wav_padding_mask, |
| | mask=None, |
| | ) |
| | return { |
| | "default": features, |
| | } |
| |
|
| |
|
