update warmup to 10

雾聪

2024-03-27 16a3cd3cfb866e99b68417e36a7c8ae613678fbf

 funasr/auto/auto_model.py

@@ -100,9 +100,7 @@
    def __init__(self, **kwargs):
        if not kwargs.get("disable_log", True):
            tables.print()
        if kwargs.get("export_model", False):
            os.environ['EXPORTING_MODEL'] = 'TRUE'
            

        model, kwargs = self.build_model(**kwargs)
        
        # if vad_model is not None, build vad model else None
@@ -157,15 +155,15 @@
            device = "cpu"
            kwargs["batch_size"] = 1
        kwargs["device"] = device
        
        if kwargs.get("ncpu", None):
            torch.set_num_threads(kwargs.get("ncpu"))

        torch.set_num_threads(kwargs.get("ncpu", 4))
        
        # build tokenizer
        tokenizer = kwargs.get("tokenizer", None)
        if tokenizer is not None:
            tokenizer_class = tables.tokenizer_classes.get(tokenizer)
            tokenizer = tokenizer_class(**kwargs["tokenizer_conf"])
            tokenizer_conf = kwargs.get("tokenizer_conf", {})
            tokenizer = tokenizer_class(**tokenizer_conf)
            kwargs["tokenizer"] = tokenizer

            kwargs["token_list"] = tokenizer.token_list if hasattr(tokenizer, "token_list") else None
@@ -478,11 +476,13 @@
               calib_num: int = 100,
               opset_version: int = 14,
               **cfg):
        os.environ['EXPORTING_MODEL'] = 'TRUE'
    
        device = cfg.get("device", "cpu")
        model = self.model.to(device=device)
        kwargs = self.kwargs
        deep_update(kwargs, cfg)
        kwargs["device"] = device
        del kwargs["model"]
        model = self.model
        model.eval()

        batch_size = 1
@@ -495,11 +495,19 @@
                export_dir = export_utils.export_onnx(
                                        model=model,
                                        data_in=data_list,
                                        quantize=quantize,
                                        fallback_num=fallback_num,
                                        calib_num=calib_num,
                                        opset_version=opset_version,
                                        **kwargs)
            else:
                export_dir = export_utils.export_torchscripts(
                                        model=model,
                                        data_in=data_list,
                                        quantize=quantize,
                                        fallback_num=fallback_num,
                                        calib_num=calib_num,
                                        opset_version=opset_version,
                                        **kwargs)

        return export_dir