Browse Source

a fix for number of threads in cluster infrastructure

runa91 3 years ago
parent
commit
17340088d4
1 changed files with 6 additions and 2 deletions
  1. 6 2
      rembg/detect.py

+ 6 - 2
rembg/detect.py

@@ -36,8 +36,12 @@ def ort_session(model_name: str) -> ort.InferenceSession:
         if hashing.hexdigest() != md5:
             with redirect_stdout(sys.stderr):
                 gdown.download(url, str(path), use_cookies=False)
-
-    return ort.InferenceSession(str(path), providers=ort.get_available_providers())
+                
+    sess_opts = ort.SessionOptions()
+    sess_opts.inter_op_num_threads = int(os.environ["OMP_NUM_THREADS"])  
+    sess_opts.intra_op_num_threads = int(os.environ["OMP_NUM_THREADS"]) 
+    
+    return ort.InferenceSession(str(path), providers=ort.get_available_providers(), sess_options=sess_opts)
 
 
 def norm_pred(d: np.ndarray) -> np.ndarray: