Databricks部署detectron2模型时出现torch模块未找到构建错误

问题根源 #

detectron2的安装过程中，其setup.py会在生成元数据阶段尝试导入torch，但此时pip尚未完成torch的安装。尽管conda_env中同时列出了torch和detectron2，但pip处理git源包时，会优先进行元数据生成，而此时torch还未被安装，导致报错。

解决方案 #

方案一：通过conda安装torch和torchvision

将torch、torchvision从pip依赖移至conda依赖项中，conda会先完成这些包的安装，再处理pip部分的依赖，确保detectron2安装时torch已存在。同时需添加pytorch通道以获取正确的包版本。

修改后的conda_env代码：

conda_env = {
    "channels": ["defaults", "pytorch"],
    "dependencies": [
        "python=3.8",
        "pip",
        "pytorch=2.0.1",
        "torchvision=0.15.2",
        "numpy",
        {
            "pip": [
                "mlflow",
                "git+https://github.com/facebookresearch/detectron2.git"
            ]
        }
    ]
}

注意：需根据detectron2版本选择兼容的torch版本，例如detectron2 v0.7+需要torch 1.9及以上版本。

方案二：使用预编译的detectron2 wheel包

预先构建与当前torch版本兼容的detectron2 wheel包，上传至DBFS后，在pip依赖中直接引用该wheel文件路径，避免从git源码编译时的依赖问题。

示例：

{
    "pip": [
        "torch",
        "mlflow",
        "torchvision",
        "numpy",
        "/dbfs/mnt/{container_name}/detectron2-0.6-cp38-cp38-linux_x86_64.whl"
    ]
}

构建wheel包时需确保与部署环境的Python版本、torch版本、CUDA版本一致。

方案三：拆分依赖安装步骤（进阶）

若必须使用pip安装所有依赖，可通过自定义MLflow环境脚本，先安装torch再安装detectron2。此方法需编写额外shell脚本，复杂度较高，不推荐作为首选方案。

替换原代码中的conda_env部分，重新注册模型即可解决安装顺序问题。

内容的提问来源于stack exchange，提问作者Stephen Dsouza