add num_threads fixes #322

This commit is contained in:
Blake Blackshear 2020-12-19 10:04:13 -06:00
parent 8d52e2635a
commit 32fb76b3d1
4 changed files with 21 additions and 10 deletions

View File

@ -200,6 +200,9 @@ detectors:
type: edgetpu
# Optional: device name as defined here: https://coral.ai/docs/edgetpu/multiple-edgetpu/#using-the-tensorflow-lite-python-api
device: usb
# Optional: num_threads value passed to the tflite.Interpreter (default: shown below)
# This value is only used for CPU types
num_threads: 3
# Optional: model configuration
model:

View File

@ -125,9 +125,9 @@ class FrigateApp():
for name, detector in self.config.detectors.items():
if detector.type == 'cpu':
self.detectors[name] = EdgeTPUProcess(name, self.detection_queue, self.detection_out_events, model_shape, tf_device='cpu')
self.detectors[name] = EdgeTPUProcess(name, self.detection_queue, self.detection_out_events, model_shape, 'cpu', detector.num_threads)
if detector.type == 'edgetpu':
self.detectors[name] = EdgeTPUProcess(name, self.detection_queue, self.detection_out_events, model_shape, tf_device=detector.device)
self.detectors[name] = EdgeTPUProcess(name, self.detection_queue, self.detection_out_events, model_shape, detector.device, detector.num_threads)
def start_detected_frames_processor(self):
self.detected_frames_processor = TrackedObjectProcessor(self.config, self.mqtt_client, self.config.mqtt.topic_prefix,

View File

@ -15,7 +15,8 @@ DETECTORS_SCHEMA = vol.Schema(
{
vol.Required(str): {
vol.Required('type', default='edgetpu'): vol.In(['cpu', 'edgetpu']),
vol.Optional('device', default='usb'): str
vol.Optional('device', default='usb'): str,
vol.Optional('num_threads', default=3): int
}
}
)
@ -264,6 +265,7 @@ class DetectorConfig():
def __init__(self, config):
self._type = config['type']
self._device = config['device']
self._num_threads = config['num_threads']
@property
def type(self):
@ -273,10 +275,15 @@ class DetectorConfig():
def device(self):
return self._device
@property
def num_threads(self):
return self._num_threads
def to_dict(self):
return {
'type': self.type,
'device': self.device
'device': self.device,
'num_threads': self.num_threads
}
class LoggerConfig():

View File

@ -43,7 +43,7 @@ class ObjectDetector(ABC):
pass
class LocalObjectDetector(ObjectDetector):
def __init__(self, tf_device=None, labels=None):
def __init__(self, tf_device=None, num_threads=3, labels=None):
self.fps = EventsPerSecond()
if labels is None:
self.labels = {}
@ -66,7 +66,7 @@ class LocalObjectDetector(ObjectDetector):
if edge_tpu_delegate is None:
self.interpreter = tflite.Interpreter(
model_path='/cpu_model.tflite')
model_path='/cpu_model.tflite', num_threads=num_threads)
else:
self.interpreter = tflite.Interpreter(
model_path='/edgetpu_model.tflite',
@ -106,7 +106,7 @@ class LocalObjectDetector(ObjectDetector):
return detections
def run_detector(name: str, detection_queue: mp.Queue, out_events: Dict[str, mp.Event], avg_speed, start, model_shape, tf_device):
def run_detector(name: str, detection_queue: mp.Queue, out_events: Dict[str, mp.Event], avg_speed, start, model_shape, tf_device, num_threads):
threading.current_thread().name = f"detector:{name}"
logger = logging.getLogger(f"detector.{name}")
logger.info(f"Starting detection process: {os.getpid()}")
@ -120,7 +120,7 @@ def run_detector(name: str, detection_queue: mp.Queue, out_events: Dict[str, mp.
signal.signal(signal.SIGINT, receiveSignal)
frame_manager = SharedMemoryFrameManager()
object_detector = LocalObjectDetector(tf_device=tf_device)
object_detector = LocalObjectDetector(tf_device=tf_device, num_threads=num_threads)
outputs = {}
for name in out_events.keys():
@ -155,7 +155,7 @@ def run_detector(name: str, detection_queue: mp.Queue, out_events: Dict[str, mp.
avg_speed.value = (avg_speed.value*9 + duration)/10
class EdgeTPUProcess():
def __init__(self, name, detection_queue, out_events, model_shape, tf_device=None):
def __init__(self, name, detection_queue, out_events, model_shape, tf_device=None, num_threads=3):
self.name = name
self.out_events = out_events
self.detection_queue = detection_queue
@ -164,6 +164,7 @@ class EdgeTPUProcess():
self.detect_process = None
self.model_shape = model_shape
self.tf_device = tf_device
self.num_threads = num_threads
self.start_or_restart()
def stop(self):
@ -179,7 +180,7 @@ class EdgeTPUProcess():
self.detection_start.value = 0.0
if (not self.detect_process is None) and self.detect_process.is_alive():
self.stop()
self.detect_process = mp.Process(target=run_detector, name=f"detector:{self.name}", args=(self.name, self.detection_queue, self.out_events, self.avg_inference_speed, self.detection_start, self.model_shape, self.tf_device))
self.detect_process = mp.Process(target=run_detector, name=f"detector:{self.name}", args=(self.name, self.detection_queue, self.out_events, self.avg_inference_speed, self.detection_start, self.model_shape, self.tf_device, self.num_threads))
self.detect_process.daemon = True
self.detect_process.start()