加载pickle文件时触发AttributeError: Can't get attribute 'test'错误

问题原因 #

pickle序列化自定义函数/类时，不会保存完整的代码逻辑，只记录函数名称和它所属的模块路径。在原笔记本中，test函数属于__main__模块（笔记本代码默认运行在该命名空间），所以同一环境加载时能找到对应定义；但新笔记本的__main__模块中没有test函数，pickle无法解析引用，因此抛出AttributeError: Can't get attribute 'test' on <module '__main__'>。

解决方案 #

方案1：加载前在当前环境定义同名函数 #

在新笔记本中先定义与原函数完全一致的test函数，再执行加载操作：

import pickle

def test(x, y):
    return x + y

filename = 'test.pkl'
loaded_model = pickle.load(open(filename, 'rb'))

方案2：将函数放到独立模块中 #

把自定义函数移到单独的Python文件（而非笔记本的__main__空间），避免依赖临时命名空间：

1. 创建my_functions.py文件，内容如下：

def test(x, y):
    return x + y

2. 原笔记本中导入并保存函数：

import pickle
from my_functions import test

filename = 'test.pkl'
pickle.dump(test, open(filename, 'wb'))

3. 新笔记本中先导入对应模块，再加载：

import pickle
import my_functions

filename = 'test.pkl'
loaded_model = pickle.load(open(filename, 'rb'))

方案3：用dill替代pickle #

dill支持序列化更多Python对象（包括__main__模块中的函数），无需额外调整代码结构：

1. 安装dill：

pip install dill

2. 原笔记本用dill保存：

import dill

def test(x, y):
    return x + y

filename = 'test.pkl'
dill.dump(test, open(filename, 'wb'))

3. 新笔记本用dill加载：

import dill

filename = 'test.pkl'
loaded_model = dill.load(open(filename, 'rb'))

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