| import pickle |
| from io import BytesIO |
| from collections import OrderedDict |
| import os |
|
|
| import torch |
|
|
|
|
| def load_pickle(path: str): |
| with open(path, "rb") as f: |
| return pickle.load(f) |
|
|
|
|
| def save_pickle(ckpt: dict, save_path: str): |
| with open(save_path, "wb") as f: |
| pickle.dump(ckpt, f) |
|
|
|
|
| def load_inputs(path: torch.serialization.FILE_LIKE, device: str, 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 export_jit_model( |
| 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 get_jit_model(model_path: str, is_half: bool, device: str, exporter): |
| jit_model_path = model_path.rstrip(".pth") |
| jit_model_path += ".half.jit" if is_half else ".jit" |
| ckpt = None |
|
|
| if os.path.exists(jit_model_path): |
| ckpt = load_pickle(jit_model_path) |
| model_device = ckpt["device"] |
| if model_device != str(device): |
| del ckpt |
| ckpt = None |
|
|
| if ckpt is None: |
| ckpt = exporter( |
| model_path=model_path, |
| mode="script", |
| inputs_path=None, |
| save_path=jit_model_path, |
| device=device, |
| is_half=is_half, |
| ) |
|
|
| return ckpt |
|
|
