浏览代码

LaTeX_OCR_rec only support MKLDNN on CPU (#2927)

Tingquan Gao 10 月之前
父节点
当前提交
aea518426e
共有 2 个文件被更改,包括 11 次插入0 次删除
  1. 10 0
      paddlex/inference/models_new/formula_recognition/predictor.py
  2. 1 0
      requirements.txt

+ 10 - 0
paddlex/inference/models_new/formula_recognition/predictor.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from ....utils import logging
 from ....utils.func_register import FuncRegister
 from ....modules.formula_recognition.model_list import MODELS
 from ...common.batch_sampler import ImageBatchSampler
@@ -65,6 +66,15 @@ class FormulaRecPredictor(BasicPredictor):
                 pre_tfs[name] = op
         pre_tfs["ToBatch"] = ToBatch()
 
+        if self.model_name in ("LaTeX_OCR_rec") and self.pp_option.device in ("cpu"):
+            import cpuinfo
+
+            if "GenuineIntel" in cpuinfo.get_cpu_info().get("vendor_id_raw", ""):
+                self.pp_option.run_mode = "mkldnn"
+                logging.warning(
+                    "Now, the `LaTeX_OCR_rec` model only support `mkldnn` mode when running on Intel CPU devices. So using `mkldnn` instead."
+                )
+
         infer = StaticInfer(
             model_dir=self.model_dir,
             model_prefix=self.MODEL_FILE_PREFIX,

+ 1 - 0
requirements.txt

@@ -1,4 +1,5 @@
 prettytable # only for benchmark
+py-cpuinfo # LaTeX_OCR_rec only support MKLDNN on CPU
 imagesize
 colorlog
 PyYAML