From 4a24df4755a592369b4574d49f4760601d126c71 Mon Sep 17 00:00:00 2001 From: Xavi Vinaixa Date: Tue, 6 Jun 2023 08:05:58 +0200 Subject: [PATCH] Incredible speed boost with TensorrtExecutionProvider #5. --- refacer.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/refacer.py b/refacer.py index c52234d..4e72732 100644 --- a/refacer.py +++ b/refacer.py @@ -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')