Files
quantulr 487b4f2e11 init
2023-09-28 15:27:08 +08:00

27 lines
678 B
Python

import os
import fastdeploy as fd
# Configurations
model_dir = "yolov5s_infer"
device = "gpu"
use_trt = True
# Prepare model
model_file = os.path.join(model_dir, "model.pdmodel")
params_file = os.path.join(model_dir, "model.pdiparams")
# Setup runtime option to select hardware, backend, etc.
option = fd.RuntimeOption()
if device.lower() == "gpu":
option.use_gpu()
if use_trt:
option.use_trt_backend()
option.set_trt_input_shape("images", [1, 3, 640, 640])
option.set_trt_cache_file("yolov5s.trt")
# Create model instance
model_instance = fd.vision.detection.YOLOv5(
model_file, params_file, runtime_option=option, model_format=fd.ModelFormat.PADDLE
)