在SageMaker中使用torch.hub加载预训练模型遇导入错误求助

问题场景 #

使用SageMaker Studio官方提供的PyTorch 2.0.0 Python 3.10 GPU优化镜像，运行在t3.medium实例上，执行以下代码时触发错误：

import torch

model = torch.hub.load("pytorch/vision", "resnet50", weights="IMAGENET1K_V2")

报错信息：

ModuleNotFoundError: No module named 'torch._custom_ops'

已通过torchvision API规避问题，以下是torch.hub触发该错误的核心原因：

原因分析 #

• 版本兼容性不匹配
  torch.hub默认会拉取pytorch/vision仓库的最新代码分支，但该分支可能已经依赖了PyTorch 2.0.0之后版本才引入的torch._custom_ops内部模块。你当前使用的PyTorch 2.0.0并未包含这个模块，导致动态加载仓库代码时触发缺失错误。而本地安装的torchvision是与PyTorch 2.0.0配套的稳定版本，自然不会出现该问题。

• torch.hub的加载机制特性
  torch.hub的工作逻辑是直接从远程仓库下载代码到本地缓存并执行，不会预先校验本地PyTorch版本与仓库代码的兼容性。当仓库代码使用了本地PyTorch未提供的内部API时，就会抛出模块缺失的异常。

• 实例与镜像的环境冲突
  你使用的t3.medium是CPU实例，但搭配的是GPU优化镜像。虽然CPU环境可以运行GPU镜像中的部分组件，但这种跨环境搭配可能导致PyTorch内部模块的依赖处理出现隐性问题，间接触发了该错误。

内容的提问来源于stack exchange，提问作者John Brown