为什么推荐使用static_dict方式保存模型

在官网教程[译]保存和加载模型中给出了多种模型使用方式,其中最常用的有

  1. 保存/加载static_dict
  2. 保存/加载完整模型

对于第一种方式,只保存训练好的模型的学习参数,但是加载时需要额外提供定义的模型结构;对于第二种方式,直接使用PyTorch的保存和加载函数即可,不过教程中也提到了第二种方式的缺陷,就是需要在调用时维护模型类文件的路径,否则会出错

之前一直使用第一种方式进行模型的读写,直到遇到了下面这个问题,才真正理解了第二种方式的缺陷

ModuleNotFoundError: No module named 'models'

使用ultralytics/yolov5的时候出现了如上错误。在网上查找了资料后发现这就是保存/加载完整模型带来的问题。参考

解析

PyTorch集成了Pickle工具进行模型的保存和加载。如果直接保存完整模型,那么附带的需要在调用时维持和模型定义文件的相对位置,否则会出现错误

解决

解决方式就是维护调用文件和模型定义文件之间的相对位置,保证Pickle能够找到模型定义文件

  1. 使用sys.path设置
import sys
sys.path.insert(0, './yolov5')
  1. 设置PYTHONPATH环境变量

相关

ultralytics/yolov5提了一个问题