Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,887 Bytes
78e32cc |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
###
# Author: Kai Li
# Date: 2021-06-17 23:08:32
# LastEditors: Please set LastEditors
# LastEditTime: 2022-05-26 18:06:22
###
import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin
def _unsqueeze_to_3d(x):
"""Normalize shape of `x` to [batch, n_chan, time]."""
if x.ndim == 1:
return x.reshape(1, 1, -1)
elif x.ndim == 2:
return x.unsqueeze(1)
else:
return x
def pad_to_appropriate_length(x, lcm):
values_to_pad = int(x.shape[-1]) % lcm
if values_to_pad:
appropriate_shape = x.shape
padded_x = torch.zeros(
list(appropriate_shape[:-1])
+ [appropriate_shape[-1] + lcm - values_to_pad],
dtype=torch.float32,
).to(x.device)
padded_x[..., : x.shape[-1]] = x
return padded_x
return x
class BaseModel(nn.Module, PyTorchModelHubMixin, repo_url="https://github.com/JusperLee/Apollo", pipeline_tag="audio-to-audio"):
def __init__(self, sample_rate, in_chan=1):
super().__init__()
self._sample_rate = sample_rate
self._in_chan = in_chan
def forward(self, *args, **kwargs):
raise NotImplementedError
def sample_rate(self,):
return self._sample_rate
@staticmethod
def load_state_dict_in_audio(model, pretrained_dict):
model_dict = model.state_dict()
update_dict = {}
for k, v in pretrained_dict.items():
if "audio_model" in k:
update_dict[k[12:]] = v
model_dict.update(update_dict)
model.load_state_dict(model_dict)
return model
@staticmethod
def from_pretrain(pretrained_model_conf_or_path, *args, **kwargs):
from . import get
conf = torch.load(
pretrained_model_conf_or_path, map_location="cpu"
) # Attempt to find the model and instantiate it.
model_class = get(conf["model_name"])
# model_class = get("Conv_TasNet")
model = model_class(*args, **kwargs)
model.load_state_dict(conf["state_dict"])
return model
def serialize(self):
import pytorch_lightning as pl # Not used in torch.hub
model_conf = dict(
model_name=self.__class__.__name__,
state_dict=self.get_state_dict(),
model_args=self.get_model_args(),
)
# Additional infos
infos = dict()
infos["software_versions"] = dict(
torch_version=torch.__version__, pytorch_lightning_version=pl.__version__,
)
model_conf["infos"] = infos
return model_conf
def get_state_dict(self):
"""In case the state dict needs to be modified before sharing the model."""
return self.state_dict()
def get_model_args(self):
"""Should return args to re-instantiate the class."""
raise NotImplementedError
|