| | |
|---|
| | | parser.add_argument("--model-name", type=str, required=True) |
|---|
| | | parser.add_argument("--export-dir", type=str, required=True) |
|---|
| | | parser.add_argument("--export", type=str2bool, default=True, help="whether to export model") |
|---|
| | | parser.add_argument("--type", type=str, default="onnx", help='["onnx", "torch"]') |
|---|
| | | parser.add_argument("--type", type=str, default="onnx", help='["onnx", "torchscript", "bladedisc"]') |
|---|
| | | parser.add_argument("--device", type=str, default="cpu", help='["cpu", "cuda"]') |
|---|
| | | parser.add_argument("--quantize", type=str2bool, default=False, help="export quantized model") |
|---|
| | | parser.add_argument("--fallback-num", type=int, default=0, help="amp fallback number") |
|---|
| | |
|---|
| | | model_file = os.path.join(model_dir, "model.onnx") |
|---|
| | | if args.quantize: |
|---|
| | | model_file = os.path.join(model_dir, "model_quant.onnx") |
|---|
| | | if args.type == "torchscript": |
|---|
| | | model_file = os.path.join(model_dir, "model.torchscript") |
|---|
| | | args.device = "cuda" |
|---|
| | | elif args.type == "bladedisc": |
|---|
| | | model_file = os.path.join(model_dir, "model_blade.torchscript") |
|---|
| | | args.device = "cuda" |
|---|
| | | if not os.path.exists(model_file): |
|---|
| | | print(".onnx is not exist, begin to export onnx") |
|---|
| | | print("model is not exist, begin to export " + model_file) |
|---|
| | | from funasr import AutoModel |
|---|
| | | |
|---|
| | | export_model = AutoModel(model=args.model_name, output_dir=output_dir) |
|---|
| | | export_model = AutoModel(model=args.model_name, output_dir=output_dir, device=args.device) |
|---|
| | | export_model.export( |
|---|
| | | quantize=args.quantize, |
|---|
| | | type=args.type, |
|---|
