Function export

funasr/utils/export_utils.py:10–63

Export. Args: model: Model instance or model name. data_in: Input data (audio samples, file paths, or text). quantize: TODO. opset_version: TODO. type: TODO. **kwargs: Additional keyword arguments.

(
    model, data_in=None, quantize: bool = False, opset_version: int = 14, type="onnx", **kwargs
)

def export(
    model, data_in=None, quantize: bool = False, opset_version: int = 14, type="onnx", **kwargs
):
    """Export.
    
        Args:
            model: Model instance or model name.
            data_in: Input data (audio samples, file paths, or text).
            quantize: TODO.
            opset_version: TODO.
            type: TODO.
            **kwargs: Additional keyword arguments.
        """
    model_scripts = model.export(**kwargs)
    export_dir = kwargs.get("output_dir", os.path.dirname(kwargs.get("init_param")))
    os.makedirs(export_dir, exist_ok=True)

    if not isinstance(model_scripts, (list, tuple)):
        model_scripts = (model_scripts,)
    for m in model_scripts:
        m.eval()
        if type == "onnx":
            _onnx(
                m,
                data_in=data_in,
                quantize=quantize,
                opset_version=opset_version,
                export_dir=export_dir,
                **kwargs,
            )
        elif type == "torchscript":
            device = "cuda" if torch.cuda.is_available() else "xpu" if torch.xpu.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
            print("Exporting torchscripts on device {}".format(device))
            _torchscripts(m, path=export_dir, device=device)
        elif type == "bladedisc":
            assert (
                torch.cuda.is_available() or torch.xpu.is_available() or torch.backends.mps.is_available()
            ), "Currently bladedisc optimization for FunASR only supports GPU"
            # bladedisc only optimizes encoder/decoder modules
            if hasattr(m, "encoder") and hasattr(m, "decoder"):
                _bladedisc_opt_for_encdec(m, path=export_dir, enable_fp16=True)
            else:
                print(f"export_dir: {export_dir}")
                _torchscripts(m, path=export_dir, device="cuda")

        elif type == "onnx_fp16":
            assert (
                torch.cuda.is_available() or torch.xpu.is_available() or torch.backends.mps.is_available()
            ), "Currently onnx_fp16 optimization for FunASR only supports GPU"

            if hasattr(m, "encoder") and hasattr(m, "decoder"):
                _onnx_opt_for_encdec(m, path=export_dir, enable_fp16=True)

    return export_dir


def _onnx(