| from io import BytesIO
|
| import pickle
|
| import time
|
| import torch
|
| from tqdm import tqdm
|
| from collections import OrderedDict
|
|
|
|
|
| def load_inputs(path, device, is_half=False):
|
| parm = torch.load(path, map_location=torch.device("cpu"))
|
| for key in parm.keys():
|
| parm[key] = parm[key].to(device)
|
| if is_half and parm[key].dtype == torch.float32:
|
| parm[key] = parm[key].half()
|
| elif not is_half and parm[key].dtype == torch.float16:
|
| parm[key] = parm[key].float()
|
| return parm
|
|
|
|
|
| def benchmark(
|
| model, inputs_path, device=torch.device("cpu"), epoch=1000, is_half=False
|
| ):
|
| parm = load_inputs(inputs_path, device, is_half)
|
| total_ts = 0.0
|
| bar = tqdm(range(epoch))
|
| for i in bar:
|
| start_time = time.perf_counter()
|
| o = model(**parm)
|
| total_ts += time.perf_counter() - start_time
|
| print(f"num_epoch: {epoch} | avg time(ms): {(total_ts*1000)/epoch}")
|
|
|
|
|
| def jit_warm_up(model, inputs_path, device=torch.device("cpu"), epoch=5, is_half=False):
|
| benchmark(model, inputs_path, device, epoch=epoch, is_half=is_half)
|
|
|
|
|
| def to_jit_model(
|
| model_path,
|
| model_type: str,
|
| mode: str = "trace",
|
| inputs_path: str = None,
|
| device=torch.device("cpu"),
|
| is_half=False,
|
| ):
|
| model = None
|
| if model_type.lower() == "synthesizer":
|
| from .get_synthesizer import get_synthesizer
|
|
|
| model, _ = get_synthesizer(model_path, device)
|
| model.forward = model.infer
|
| elif model_type.lower() == "rmvpe":
|
| from .get_rmvpe import get_rmvpe
|
|
|
| model = get_rmvpe(model_path, device)
|
| elif model_type.lower() == "hubert":
|
| from .get_hubert import get_hubert_model
|
|
|
| model = get_hubert_model(model_path, device)
|
| model.forward = model.infer
|
| else:
|
| raise ValueError(f"No model type named {model_type}")
|
| model = model.eval()
|
| model = model.half() if is_half else model.float()
|
| if mode == "trace":
|
| assert not inputs_path
|
| inputs = load_inputs(inputs_path, device, is_half)
|
| model_jit = torch.jit.trace(model, example_kwarg_inputs=inputs)
|
| elif mode == "script":
|
| model_jit = torch.jit.script(model)
|
| model_jit.to(device)
|
| model_jit = model_jit.half() if is_half else model_jit.float()
|
|
|
| return (model, model_jit)
|
|
|
|
|
| def export(
|
| model: torch.nn.Module,
|
| mode: str = "trace",
|
| inputs: dict = None,
|
| device=torch.device("cpu"),
|
| is_half: bool = False,
|
| ) -> dict:
|
| model = model.half() if is_half else model.float()
|
| model.eval()
|
| if mode == "trace":
|
| assert inputs is not None
|
| model_jit = torch.jit.trace(model, example_kwarg_inputs=inputs)
|
| elif mode == "script":
|
| model_jit = torch.jit.script(model)
|
| model_jit.to(device)
|
| model_jit = model_jit.half() if is_half else model_jit.float()
|
| buffer = BytesIO()
|
|
|
| torch.jit.save(model_jit, buffer)
|
| del model_jit
|
| cpt = OrderedDict()
|
| cpt["model"] = buffer.getvalue()
|
| cpt["is_half"] = is_half
|
| return cpt
|
|
|
|
|
| def load(path: str):
|
| with open(path, "rb") as f:
|
| return pickle.load(f)
|
|
|
|
|
| def save(ckpt: dict, save_path: str):
|
| with open(save_path, "wb") as f:
|
| pickle.dump(ckpt, f)
|
|
|
|
|
| def rmvpe_jit_export(
|
| model_path: str,
|
| mode: str = "script",
|
| inputs_path: str = None,
|
| save_path: str = None,
|
| device=torch.device("cpu"),
|
| is_half=False,
|
| ):
|
| if not save_path:
|
| save_path = model_path.rstrip(".pth")
|
| save_path += ".half.jit" if is_half else ".jit"
|
| if "cuda" in str(device) and ":" not in str(device):
|
| device = torch.device("cuda:0")
|
| from .get_rmvpe import get_rmvpe
|
|
|
| model = get_rmvpe(model_path, device)
|
| inputs = None
|
| if mode == "trace":
|
| inputs = load_inputs(inputs_path, device, is_half)
|
| ckpt = export(model, mode, inputs, device, is_half)
|
| ckpt["device"] = str(device)
|
| save(ckpt, save_path)
|
| return ckpt
|
|
|
|
|
| def synthesizer_jit_export(
|
| model_path: str,
|
| mode: str = "script",
|
| inputs_path: str = None,
|
| save_path: str = None,
|
| device=torch.device("cpu"),
|
| is_half=False,
|
| ):
|
| if not save_path:
|
| save_path = model_path.rstrip(".pth")
|
| save_path += ".half.jit" if is_half else ".jit"
|
| if "cuda" in str(device) and ":" not in str(device):
|
| device = torch.device("cuda:0")
|
| from .get_synthesizer import get_synthesizer
|
|
|
| model, cpt = get_synthesizer(model_path, device)
|
| assert isinstance(cpt, dict)
|
| model.forward = model.infer
|
| inputs = None
|
| if mode == "trace":
|
| inputs = load_inputs(inputs_path, device, is_half)
|
| ckpt = export(model, mode, inputs, device, is_half)
|
| cpt.pop("weight")
|
| cpt["model"] = ckpt["model"]
|
| cpt["device"] = device
|
| save(cpt, save_path)
|
| return cpt
|
|
|
