TensorrtExecutionProvider

This commit is contained in:
Xavi Vinaixa
2023-06-06 08:14:28 +02:00
parent a8d2941293
commit 8008e2e766

View File

@@ -49,16 +49,17 @@ class Refacer:
self.use_num_cpus = mp.cpu_count()-1 self.use_num_cpus = mp.cpu_count()-1
print(f"CoreML mode with providers {self.providers}") print(f"CoreML mode with providers {self.providers}")
self.sess_options.intra_op_num_threads = int(self.use_num_cpus/2) self.sess_options.intra_op_num_threads = int(self.use_num_cpus/2)
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: elif 'TensorrtExecutionProvider' in self.providers:
self.mode = RefacerMode.TENSORRT self.mode = RefacerMode.TENSORRT
self.use_num_cpus = 1 self.use_num_cpus = 1
self.sess_options.intra_op_num_threads = 1 self.sess_options.intra_op_num_threads = 1
print(f"TENSORT mode with providers {self.providers}") print(f"TENSORT mode with providers {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}")
def __init_apps(self): def __init_apps(self):
assets_dir = ensure_available('models', 'buffalo_l', root='~/.insightface') assets_dir = ensure_available('models', 'buffalo_l', root='~/.insightface')