From 2e769fb36ce88dabfa984e8b81e8cb1c90799c95 Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期五, 07 四月 2023 15:54:09 +0800
Subject: [PATCH] Merge branch 'main' into dev_cmz2
---
funasr/export/export_model.py | 8 ++++++++
1 files changed, 8 insertions(+), 0 deletions(-)
diff --git a/funasr/export/export_model.py b/funasr/export/export_model.py
index 444ccf4..b69eeee 100644
--- a/funasr/export/export_model.py
+++ b/funasr/export/export_model.py
@@ -19,6 +19,7 @@
self,
cache_dir: Union[Path, str] = None,
onnx: bool = True,
+ device: str = "cpu",
quant: bool = True,
fallback_num: int = 0,
audio_in: str = None,
@@ -36,6 +37,7 @@
)
print("output dir: {}".format(self.cache_dir))
self.onnx = onnx
+ self.device = device
self.quant = quant
self.fallback_num = fallback_num
self.frontend = None
@@ -111,6 +113,10 @@
dummy_input = model.get_dummy_inputs(enc_size)
else:
dummy_input = model.get_dummy_inputs()
+
+ if self.device == 'cuda':
+ model = model.cuda()
+ dummy_input = tuple([i.cuda() for i in dummy_input])
# model_script = torch.jit.script(model)
model_script = torch.jit.trace(model, dummy_input)
@@ -260,6 +266,7 @@
parser.add_argument('--model-name', type=str, required=True)
parser.add_argument('--export-dir', type=str, required=True)
parser.add_argument('--type', type=str, default='onnx', help='["onnx", "torch"]')
+ 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')
parser.add_argument('--audio_in', type=str, default=None, help='["wav", "wav.scp"]')
@@ -269,6 +276,7 @@
export_model = ModelExport(
cache_dir=args.export_dir,
onnx=args.type == 'onnx',
+ device=args.device,
quant=args.quantize,
fallback_num=args.fallback_num,
audio_in=args.audio_in,
--
Gitblit v1.9.1