使用D2Go训练自定义数据集遇ModuleNotFoundError求助

问题场景 #

我尝试用D2Go训练自定义数据集，按官方指引执行了以下安装命令：

!git clone https://github.com/facebookresearch/d2go
!python -m pip install ./d2go

但在Google Colab中运行from d2go.model_zoo import model_zoo时，抛出如下错误：

ModuleNotFoundError Traceback (most recent call last)
in
----> 1 from d2go.model_zoo import model_zoo

5 frames
/content/d2go/d2go/trainer/fsdp.py in
14 from d2go.trainer.helper import parse_precision_from_string
15 from detectron2.utils.registry import Registry
---> 16 from torch.ao.pruning import fqn_to_module
17 from torch.cuda.amp import GradScaler
18 from torch.distributed.fsdp.fully_sharded_data_parallel import (

ModuleNotFoundError: No module named 'torch.ao.pruning'

解决方案 #

问题原因 #

torch.ao.pruning模块是PyTorch 1.10及以上版本才新增的功能，Google Colab默认预装的PyTorch版本可能低于该要求，导致找不到模块。

修复步骤 #

1. 卸载当前低版本PyTorch：

!pip uninstall -y torch torchvision torchaudio

2. 安装符合版本要求的PyTorch（以适配Colab的CUDA 11.8为例）：

!pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118

3. 重新安装D2Go及其完整依赖：

!pip install -e ./d2go[all]

4. 重新验证导入：

from d2go.model_zoo import model_zoo
print("D2Go模型库导入成功")

备选方案 #

如果上述版本仍有问题，可以查看D2Go仓库根目录的requirements.txt文件，确认官方指定的PyTorch版本，安装对应版本后再重新安装D2Go。

内容的提问来源于stack exchange，提问作者c.dipu0