admin/LC_NeoRefacer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Incredible speed boost with TensorrtExecutionProvider #5.

Browse Source
This commit is contained in:
Xavi Vinaixa
2023-06-06 08:05:58 +02:00
parent 7e628a3b83
commit 4a24df4755
1 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
refacer.py
View File

@@ -20,7 +20,7 @@ from insightface.app.common import Face
from insightface.utils.storage import ensure_available
class RefacerMode(Enum):
CPU, CUDA, COREML = range(1, 4)
CPU, CUDA, COREML, TENSORRT = range(1, 4)
class Refacer:
def __init__(self,force_cpu=False):
@@ -49,11 +49,17 @@ class Refacer:
self.use_num_cpus = mp.cpu_count()-1
print(f"CoreML mode with providers {self.providers}")
self.sess_options.intra_op_num_threads = int(self.use_num_cpus/2)
elif 'CUDAExecutionProvider' in self.providers or 'TensorrtExecutionProvider' in self.providers:
elif 'CUDAExecutionProvider' in self.providers:
self.mode = RefacerMode.CUDA
self.use_num_cpus = 1
self.sess_options.intra_op_num_threads = 1
print(f"CUDA mode with providers {self.providers}")
elif 'TensorrtExecutionProvider' in self.providers:
self.mode = RefacerMode.TENSORRT
self.use_num_cpus = 1
self.sess_options.intra_op_num_threads = 1
print(f"TENSORT mode with providers {self.providers}")
def __init_apps(self):
assets_dir = ensure_available('models', 'buffalo_l', root='~/.insightface')