调用yolov5训好的本地模型pt

1、新建一个文件夹pdf2picture

2、将yolov5项目中的模型yolov5s-flash.pt复制到该目录下

3、必须在yolov5项目中的hubconf.py中定义可调用(入口点)的名称：

4、新建的文件夹下建议一个pdf_yolov5_to_flash.py：

import torch

def yolov5_model():
    # GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    # Model
    model = torch.hub.load('../yolov5',   
                          'yolov5s_flash',
                           pretrained=True,
                           source='local')  # or yolov5m, yolov5l, yolov5x, custom
    model = model.to(device)
    return model

model = yolov5_model()

注意：

    yolov5s-flash.pt也必须在yolov5项目的根目录下

    model目录下的yaml名称要与pt文件名、hubconf.py入口点名称一致，这里统一是“yolov5s_flash”