robmarkcole / yolov5-fastapi

FastAPI app exposing yolov5 object detection
Apache License 2.0
10 stars 2 forks source link

load self-trained yolov5 model #6

Closed Amanda-Barbara closed 3 years ago

Amanda-Barbara commented 3 years ago

@robmarkcole Hi, I want to load my self-trained yolov5 model and to use GPU mode, where I need to modify? thanks.

robmarkcole commented 3 years ago

You need to modify here https://github.com/robmarkcole/yolov5-fastapi/blob/78597d88269fe415f96a871548f2a7b5f8961477/celery_tasks/yolo.py#L9

Amanda-Barbara commented 3 years ago

@robmarkcole thanks.