allow setting specific edgetpu in config

This commit is contained in:
Blake Blackshear 2020-09-07 12:49:47 -05:00
parent f64320a464
commit 2f758af097
3 changed files with 22 additions and 8 deletions

View File

@ -1,5 +1,11 @@
web_port: 5000 web_port: 5000
################
## Tell frigate to look for a specific EdgeTPU device. Useful if you want to run multiple instances of frigate
## on the same machine with multiple EdgeTPUs. https://coral.ai/docs/edgetpu/multiple-edgetpu/#using-the-tensorflow-lite-python-api
################
tensorflow_device: usb
mqtt: mqtt:
host: mqtt.server.com host: mqtt.server.com
topic_prefix: frigate topic_prefix: frigate

View File

@ -61,6 +61,7 @@ GLOBAL_OBJECT_CONFIG = CONFIG.get('objects', {})
WEB_PORT = CONFIG.get('web_port', 5000) WEB_PORT = CONFIG.get('web_port', 5000)
DEBUG = (CONFIG.get('debug', '0') == '1') DEBUG = (CONFIG.get('debug', '0') == '1')
TENSORFLOW_DEVICE = CONFIG.get('tensorflow_device')
def start_plasma_store(): def start_plasma_store():
plasma_cmd = ['plasma_store', '-m', '400000000', '-s', '/tmp/plasma'] plasma_cmd = ['plasma_store', '-m', '400000000', '-s', '/tmp/plasma']
@ -190,7 +191,7 @@ def main():
event_queue = mp.Queue() event_queue = mp.Queue()
# Start the shared tflite process # Start the shared tflite process
tflite_process = EdgeTPUProcess() tflite_process = EdgeTPUProcess(TENSORFLOW_DEVICE)
# start the camera processes # start the camera processes
camera_processes = {} camera_processes = {}

View File

@ -34,18 +34,24 @@ class ObjectDetector(ABC):
pass pass
class LocalObjectDetector(ObjectDetector): class LocalObjectDetector(ObjectDetector):
def __init__(self, labels=None): def __init__(self, tf_device=None, labels=None):
if labels is None: if labels is None:
self.labels = {} self.labels = {}
else: else:
self.labels = load_labels(labels) self.labels = load_labels(labels)
device_config = {"device": "usb"}
if not tf_device is None:
device_config = {"device": tf_device}
edge_tpu_delegate = None edge_tpu_delegate = None
try: try:
edge_tpu_delegate = load_delegate('libedgetpu.so.1.0', {"device": "usb"}) print(f"Attempting to load TPU as {device_config['device']}")
print("USB TPU found") edge_tpu_delegate = load_delegate('libedgetpu.so.1.0', device_config)
print("TPU found")
except ValueError: except ValueError:
try: try:
print(f"Attempting to load TPU as pci:0")
edge_tpu_delegate = load_delegate('libedgetpu.so.1.0', {"device": "pci:0"}) edge_tpu_delegate = load_delegate('libedgetpu.so.1.0', {"device": "pci:0"})
print("PCIe TPU found") print("PCIe TPU found")
except ValueError: except ValueError:
@ -92,11 +98,11 @@ class LocalObjectDetector(ObjectDetector):
return detections return detections
def run_detector(detection_queue, avg_speed, start): def run_detector(detection_queue, avg_speed, start, tf_device):
print(f"Starting detection process: {os.getpid()}") print(f"Starting detection process: {os.getpid()}")
listen() listen()
plasma_client = plasma.connect("/tmp/plasma") plasma_client = plasma.connect("/tmp/plasma")
object_detector = LocalObjectDetector() object_detector = LocalObjectDetector(tf_device=tf_device)
while True: while True:
object_id_str = detection_queue.get() object_id_str = detection_queue.get()
@ -117,11 +123,12 @@ def run_detector(detection_queue, avg_speed, start):
avg_speed.value = (avg_speed.value*9 + duration)/10 avg_speed.value = (avg_speed.value*9 + duration)/10
class EdgeTPUProcess(): class EdgeTPUProcess():
def __init__(self): def __init__(self, tf_device=None):
self.detection_queue = mp.Queue() self.detection_queue = mp.Queue()
self.avg_inference_speed = mp.Value('d', 0.01) self.avg_inference_speed = mp.Value('d', 0.01)
self.detection_start = mp.Value('d', 0.0) self.detection_start = mp.Value('d', 0.0)
self.detect_process = None self.detect_process = None
self.tf_device = tf_device
self.start_or_restart() self.start_or_restart()
def start_or_restart(self): def start_or_restart(self):
@ -134,7 +141,7 @@ class EdgeTPUProcess():
print("Detection process didnt exit. Force killing...") print("Detection process didnt exit. Force killing...")
self.detect_process.kill() self.detect_process.kill()
self.detect_process.join() self.detect_process.join()
self.detect_process = mp.Process(target=run_detector, args=(self.detection_queue, self.avg_inference_speed, self.detection_start)) self.detect_process = mp.Process(target=run_detector, args=(self.detection_queue, self.avg_inference_speed, self.detection_start, self.tf_device))
self.detect_process.daemon = True self.detect_process.daemon = True
self.detect_process.start() self.detect_process.start()