export

游雁

2023-06-20 43ad2c35634a3ed2a7a46bd7e3afd147934b1c48

export

 funasr/export/export_model.py

@@ -229,7 +229,7 @@
        # model_script = torch.jit.script(model)
        model_script = model #torch.jit.trace(model)
        model_path = os.path.join(path, f'{model.model_name}.onnx')

        if not os.path.exists(model_path):
        torch.onnx.export(
            model_script,
            dummy_input,
@@ -245,6 +245,7 @@
            from onnxruntime.quantization import QuantType, quantize_dynamic
            import onnx
            quant_model_path = os.path.join(path, f'{model.model_name}_quant.onnx')
            if not os.path.exists(quant_model_path):
            onnx_model = onnx.load(model_path)
            nodes = [n.name for n in onnx_model.graph.node]
            nodes_to_exclude = [m for m in nodes if 'output' in m]