将 PyTorch 模型加载到我的 Flask 应用程序中，会抛出此错误 - ModuleNotFoundError: No module named 'models'

Question

该项目旨在从图像中检测对象并生成一个音频文件，其中音频播放对象的真实内容。例如，如果对象包含一辆汽车，则音频文件包含一个音频“CAR”。但这里的问题是，当我尝试运行 Flask 应用程序时，我遇到了这个错误：

PS E:\ObjRec> python app.py
Loading YOLOv5 model...
Traceback (most recent call last):
  File "E:\ObjRec\app.py", line 11, in <module>
    model = torch.load('best.pt')
            ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Admin\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\torch\serialization.py", line 1026, in load
    return _load(opened_zipfile,
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Admin\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\torch\serialization.py", line 1438, in _load
    result = unpickler.load()
             ^^^^^^^^^^^^^^^^
  File "C:\Users\Admin\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\torch\serialization.py", line 1431, in find_class
    return super().find_class(mod_name, name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'models'

这是我的

app.py

：

from flask import Flask, request, jsonify
import torch
from torchvision import transforms
from PIL import Image
import io

app = Flask(__name__)

print('Loading YOLOv5 model...')
model = torch.load('best.pt')
model.eval()

preprocess = transforms.Compose([
    transforms.Resize((416, 416)),
    transforms.ToTensor(),
])

@app.route('/detect', methods=['POST'])
def detect_objects():
    if 'image' not in request.files:
        return jsonify({'error': 'No image uploaded'}), 400
    image_file = request.files['image']
    image_bytes = image_file.read()
    image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
    image = preprocess(image)
    results = model(image.unsqueeze(0)) 
    labels = results.names
    
    return jsonify({'labels': labels}), 200

if __name__ == '__main__':
    app.run(debug=True)

如果不将模型加载到项目中，我就无法继续进行。我也不确定我遵循的程序是否正确。

这是实际的文件路径。

我尝试

pip install models

并再次遇到元数据生成失败错误：

PS E:\ObjRec> pip install models
Collecting models
  Using cached models-0.9.3.tar.gz (16 kB)
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [8 lines of output]
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
        File "<pip-setuptools-caller>", line 34, in <module>
        File "C:\Users\Admin\AppData\Local\Temp\pip-install-8yqo_00s\models_16df2befc82545e7a9071610e1043cde\setup.py", line 25, in <module>
          import models
        File "C:\Users\Admin\AppData\Local\Temp\pip-install-8yqo_00s\models_16df2befc82545e7a9071610e1043cde\models\__init__.py", line 23, in <module>
          from base import *
      ModuleNotFoundError: No module named 'base'
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.

不知道如何从这里继续前进！

Answer 1

问题在于如何加载模型以及如何保存模型。 序列化模型（'best.pt'）绑定到训练期间使用的特定类和确切的目录结构（来自 models.Yolo import YOLOv5），并且只能加载到训练时的同一位置。

解决这个问题的最佳方法：

1. 在包含训练模型脚本的目录中（以便导入与训练期间使用的导入相同）加载有界模型：

   模型 = torch.load(PATH.pt)

2. 从加载的模型中提取state_dict()并保存。

   torch.save(model.state_dict(), PATH_UNBOUNDED.pt)

3. 在不同的项目/位置加载无界模型

   model = YOLO()#用YOLO定义实例化类 检查点 = torch.load(PATH_UNBOUNDED.pt) model.load_state_dict(checkpoint)#填充权重