使用FastAI load_learner()时遇__builtin__模块找不到错误

问题背景 #

• 为减少代码量，希望使用FastAI的load_learner()加载预存模型，而非每次重新创建learner实例
• 曾用dill替代pickle保存模型，怀疑存在兼容性问题
• 调用load_learner()时触发报错：ModuleNotFoundError: No module named '__builtin__'，已知Python 3中__builtins__已更名为builtins

环境版本 #

torch==1.7.1
fastai==2.7.7
fastcore==1.5.6
torchvision==0.8.2
Python 3.9

相关代码（predict.py） #

if __name__ == '__main__':
    file_path = project_params.file_path
    learner = project_params.learner
    model = load_learner(file_path.model, cpu=learner.cpu)  # 报错行

配置文件（params.yaml） #

learner:
  cpu: False
# ...
file_path:
  model: ./project_model/data/learner.pkl

完整报错日志 #

(venv) me@ubuntu-pcs:~/PycharmProjects/project$ python project_model/predict.py
'foo/data/bar.dvc' didn't change, skipping    
Running stage 'predict':
> python foo/predict.py
/home/me/miniconda3/envs/venv/lib/python3.9/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
Traceback (most recent call last):
  File "/home/me/PycharmProjects/project/foo/predict.py", line 35, in <module>
    model = load_learner(file_path.model, cpu=learner.cpu)
  File "/home/me/miniconda3/envs/venv/lib/python3.9/site-packages/fastai/learner.py", line 414, in load_learner
    try: res = torch.load(fname, map_location=map_loc, pickle_module=pickle_module)
  File "/home/me/miniconda3/envs/venv/lib/python3.9/site-packages/torch/serialization.py", line 595, in load
    return _legacy_load(opened_file, map_location=pickle_module, **pickle_load_args)
  File "/home/me/miniconda3/envs/venv/lib/python3.9/site-packages/torch/serialization.py", line 764, in _legacy_load
    magic_number = pickle_module.load(f, **pickle_load_args)
ModuleNotFoundError: No module named '__builtin__'
ERROR: failed to reproduce 'predict': failed to run: python foo/predict.py, exited with 1

解决方案 #

方案1：修复dill序列化的兼容问题 #

如果模型是用dill保存的，加载时需指定pickle模块为dill，并添加Python 3的兼容补丁：

import builtins
import dill
from fastai.learner import load_learner

# 给dill添加__builtin__别名，适配Python3的命名变更
dill._dill._builtin = builtins

# 加载模型时指定pickle_module参数
model = load_learner(file_path.model, cpu=learner.cpu, pickle_module=dill)

方案2：使用FastAI原生方法重新保存模型 #

若上述方案无效，建议放弃dill，改用FastAI官方的export()方法保存模型：

# 在训练代码中执行导出
learn.export("learner.pkl")

之后直接用load_learner()加载即可，无需额外配置。

方案3：确保环境版本一致性 #

检查保存模型和加载模型的Python、FastAI、PyTorch版本完全一致，避免因版本差异导致的序列化不兼容问题。

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