如何正确加载pytorch pth模型并与另一个模型结合？答案

【问题标题】：How to load pytorch pth model correctly and combine with another model?如何正确加载pytorch pth模型并与另一个模型结合？
【发布时间】：2021-10-11 18:10:07
【问题描述】：

我正在寻找与 pytorch 一起使用的 OCR 预训练模型。我试过https://github.com/clovaai/CRAFT-pytorch，但它在 pytorch hub 上不受支持。而且我无法加载 pth 模型，因为它只有权重。如何加载模型？

我的第一个模型是使用自定义数据训练的 yolov5 模型，因此它应该裁剪图像并将其发送到下一个模型。下一个模型应该是 OCR，主要是数字识别。但是我不能运行craft-pytorch

model = torch.hub.load('.', 'custom', path='runs/train/exp2/weights/best.pt', source='local', force_reload=True)

# It throws an error with pytorch hub
model_ocr = torch.hub.load('clovaai/CRAFT-pytorch', 'craft_mlt_25k.pth')
# Tried with torch load, but pth files have only weights not model
ocr_model = torch.load('runs/ocr/craft_mlt_25k.pth')

cap = cv2.VideoCapture('../Dataset/test/09-10.mp4')

while(cap.isOpened()):
    ret, frame = cap.read()
    results = model(frame)
    crops = results.crop(save=False)    
    for crop in crops:
        if 'number' in crop['label']:
            ocr_result = model_ocr(crop['im'])
            ocr_crop = ocr_result.crop(save=False)

【问题讨论】：

标签： python pytorch ocr

【解决方案1】：

如果你有模型的源代码，你可以创建它并加载它的权重，但正如我理解的问题，你希望在没有源代码的情况下这样做。

为此，您只需致电torch.load：

state_dict = torch.load('model_file.pth')

state_dict 应该只是一个（类型的）字典。您可以根据需要检查和挑选它：

print(state_dict.keys())

with torch.no_grad():
    my_model.some_weight[:] = state_dict['some_other_name']

（请注意，在不受信任的数据上调用 torch.load 是不安全的）

【讨论】：