From 2a04e1a12a422941b55e794df9d49fcc312b2e0b Mon Sep 17 00:00:00 2001
From: LDOUBLEV <liuvv0203@outlook.com>
Date: Tue, 22 Dec 2020 15:57:21 +0800
Subject: [PATCH] add export to limit GPU-memory-usage

---
 tools/infer/predict_cls.py    | 2 ++
 tools/infer/predict_det.py    | 2 ++
 tools/infer/predict_rec.py    | 2 ++
 tools/infer/predict_system.py | 2 ++
 tools/infer/utility.py        | 5 +++--
 tools/infer_cls.py            | 2 ++
 tools/infer_det.py            | 2 ++
 tools/infer_rec.py            | 2 ++
 8 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/tools/infer/predict_cls.py b/tools/infer/predict_cls.py
index 6604d3cc..074172cc 100755
--- a/tools/infer/predict_cls.py
+++ b/tools/infer/predict_cls.py
@@ -18,6 +18,8 @@ __dir__ = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(__dir__)
 sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
 
+os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
+
 import cv2
 import copy
 import numpy as np
diff --git a/tools/infer/predict_det.py b/tools/infer/predict_det.py
index e677eb24..077692af 100755
--- a/tools/infer/predict_det.py
+++ b/tools/infer/predict_det.py
@@ -18,6 +18,8 @@ __dir__ = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(__dir__)
 sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
 
+os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
+
 import cv2
 import numpy as np
 import time
diff --git a/tools/infer/predict_rec.py b/tools/infer/predict_rec.py
index 89d04cff..974fdbb6 100755
--- a/tools/infer/predict_rec.py
+++ b/tools/infer/predict_rec.py
@@ -18,6 +18,8 @@ __dir__ = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(__dir__)
 sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
 
+os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
+
 import cv2
 import numpy as np
 import math
diff --git a/tools/infer/predict_system.py b/tools/infer/predict_system.py
index 07dfc216..8c4f9214 100755
--- a/tools/infer/predict_system.py
+++ b/tools/infer/predict_system.py
@@ -18,6 +18,8 @@ __dir__ = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(__dir__)
 sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
 
+os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
+
 import cv2
 import copy
 import numpy as np
diff --git a/tools/infer/utility.py b/tools/infer/utility.py
index 39c045de..966fa3cc 100755
--- a/tools/infer/utility.py
+++ b/tools/infer/utility.py
@@ -33,7 +33,7 @@ def parse_args():
     parser.add_argument("--ir_optim", type=str2bool, default=True)
     parser.add_argument("--use_tensorrt", type=str2bool, default=False)
     parser.add_argument("--use_fp16", type=str2bool, default=False)
-    parser.add_argument("--gpu_mem", type=int, default=8000)
+    parser.add_argument("--gpu_mem", type=int, default=500)
 
     # params for text detector
     parser.add_argument("--image_dir", type=str)
@@ -62,7 +62,7 @@ def parse_args():
     parser.add_argument("--rec_model_dir", type=str)
     parser.add_argument("--rec_image_shape", type=str, default="3, 32, 320")
     parser.add_argument("--rec_char_type", type=str, default='ch')
-    parser.add_argument("--rec_batch_num", type=int, default=1)
+    parser.add_argument("--rec_batch_num", type=int, default=6)
     parser.add_argument("--max_text_length", type=int, default=25)
     parser.add_argument(
         "--rec_char_dict_path",
@@ -123,6 +123,7 @@ def create_predictor(args, mode, logger):
             # cache 10 different shapes for mkldnn to avoid memory leak
             config.set_mkldnn_cache_capacity(10)
             config.enable_mkldnn()
+            args.rec_batch_num = 1
 
     # config.enable_memory_optim()
     config.disable_glog_info()
diff --git a/tools/infer_cls.py b/tools/infer_cls.py
index 85e11ac3..49696482 100755
--- a/tools/infer_cls.py
+++ b/tools/infer_cls.py
@@ -25,6 +25,8 @@ __dir__ = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(__dir__)
 sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
 
+os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
+
 import paddle
 
 from ppocr.data import create_operators, transform
diff --git a/tools/infer_det.py b/tools/infer_det.py
index d1b1b752..d890970e 100755
--- a/tools/infer_det.py
+++ b/tools/infer_det.py
@@ -25,6 +25,8 @@ __dir__ = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(__dir__)
 sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
 
+os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
+
 import cv2
 import json
 import paddle
diff --git a/tools/infer_rec.py b/tools/infer_rec.py
index e3e85b5d..7e4b0811 100755
--- a/tools/infer_rec.py
+++ b/tools/infer_rec.py
@@ -25,6 +25,8 @@ __dir__ = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(__dir__)
 sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))
 
+os.environ["FLAGS_allocator_strategy"] = 'auto_growth'
+
 import paddle
 
 from ppocr.data import create_operators, transform
-- 
GitLab