pickle反序列化时的ImportError

pickle常用来序列号和反序列化Python对象，尽管从软件工程角度看pickle不是一个好的选择，但架不住其在ML生态的广泛使用。当然，本文并非吐槽pickle，毕竟好与不好都得用，本文尝试解决pickle反序列化时可能存在的一个异常。

如文档所述， pickle 库是可以保存和恢复类实例对象的，但是类定义必须是可导入的，且类的定义须与类对象存储时位于相同的模块（及保持相同的目录结构）中。

pickle can save and restore class instances transparently, however the class definition must be importable and live in the same module as when the object was stored.

那按照文档的描述，如果在保存类对象后修改了类定义所在模块的目录结构，那保存的类对象文件就没法被加载了吗？不应该吧，毕竟我类定义和类对象序列化后的二进制文件都还在，凭什么就不能被加载了呢？

本文尝试复现该问题并给出解决方案。

问题复现与分析

先来描述一下问题。假设我目前的代码结构如下：

proj
    │  cat.py
    │  main.py

在cat.py中定义了Cat类，然后在main.py中实例化了一个Cat类，且通过joblib.dump将实例化对象保存为了二进制文件my_cat.pkl。代码如下：

# cat.py
class Cat:
    def __init__(self, age: int, weight: float, color: str) -> None:
        self.age = age
        self.weight = weight
        self.color = color
    def speak(self):
        print("Meow")
    def __str__(self) -> str:
        return f"Cat: age={self.age}, weight={self.weight}, color={self.color}"import joblib
        
        
# main.py     
from cat import Cat
my_cat = Cat(3, 5.5, "black")
joblib.dump(my_cat, "my_cat.pkl")

此时，二进制文件my_cat.pkl保存至项目根目录。

接下来对目录结构进行些许调整，如下：

proj
    │  main.py
    │  my_cat.pkl
    └─pet
      │ cat.py
      │ __init__.py

且修改main.py代码，以加载之前保存的Cat对象。

import joblib
my_cat = joblib.load("my_cat.pkl")

执行main.py之后，有如下报错信息：

Traceback (most recent call last):
  File "C:\Users\xxxxxx\Desktop\proj\main.py", line 5, in <module>    my_cat = joblib.load("my_cat.pkl")
  File "D:\miniconda3\envs\dl_torch\lib\site-packages\joblib\numpy_pickle.py", line 658, in load
    obj = _unpickle(fobj, filename, mmap_mode)
  File "D:\miniconda3\envs\dl_torch\lib\site-packages\joblib\numpy_pickle.py", line 577, in _unpickle
    obj = unpickler.load()
  File "D:\miniconda3\envs\dl_torch\lib\pickle.py", line 1212, in load
    dispatch[key[0]](self)
  File "D:\miniconda3\envs\dl_torch\lib\pickle.py", line 1537, in load_stack_global
    self.append(self.find_class(module, name))
  File "D:\miniconda3\envs\dl_torch\lib\pickle.py", line 1579, in find_class
    __import__(module, level=0)
ModuleNotFoundError: No module named 'cat'

注意此处是找不到名为cat的module，而不是找不到类定义。那它是怎么找的呢？一个二进制文件在load的时候要如何知道我原先的类定义在哪个模块？又是如何知道我这个模块的位置呢？

packle在dump类对象时候，顺便保存了类定义的引用，即保存了import 路径，然后在load文件的时候，会以相同的方式尝试import这个模块。到此便能理解为什么pickle文档所说的 the class definition must be importable and live in the same module as when the object was stored. 这句话了。

综上，该问题的本质原因就是pickle在dump类对象时保存了类定义的引用（本例中为 cat.Cat），然后当你改变类定义所在模块的目录结构后（本例修改为了pet.cat.Cat），pkl二进制文件中的类定义的引用没用也没法被同步被更改，然后在load的时候，依旧以同样的路径查找类定义，当然找不到了，甚至连模块cat都找不到，因此会报错ModuleNotFoundError: No module named 'cat'

解决思路

pickle找不到那就帮它找咯！上文讲了，在load的时候，是从pkl二进制文件中得到类定义的路径，然后再import该路径。既然如此，便有如下两个思路：

• 直接打开该二进制文件，然后修改其中的类定义所在模块路径为当前模块路径不就可以吗？理论上是可以的，但由于序列化的对象为二进制文件，操作起来有风险且有难度，因此不建议这样操作。
• 在load时，代码按照二进制文件中保存的模块路径查找模块的时候，进行一次拦截，用当前的模块的路径来替换，使之能成功上位。下文着重实操此法。

要知道怎么拦截，需要先明白Python如何import模块。此时了解一下Python中sys.modules的作用。

sys.modules

可见，sys.modules保存了当前解释器已经加载好的模块名称和模块。然后，当需要某个模块时，会首先在sys.modules中查找，查看是否已经被import了，如果在sys.modules找到了，那便不会重复import。

明白了import机制，解决方法就很自然了：

1. 首先按照当前的文件目录import类定义所在的模块。本例中为import pet
2. 然后修改sys.modules，追加旧模块与新模块的映射。本例为sys.modules[“cat”] = pet.cat
3. 大功告成，直接load文件!

解释如下：在load二进制文件之前，便将旧的模块名称存放到sys.modules中，load文件时，当解释器看到二进制文件中的模块路径之后，首先会在sys.modules查找，当发现已经存在了，那便不会按照旧的路径再查找一遍了。

具体操作如下：

修改main.py如下

import sys
import joblib
import pet
sys.modules["cat"] = pet.cat
my_cat = joblib.load("my_cat.pkl")
print(type(my_cat))

此时，便可正常load文件my_cat.pkl，且打印”<class ‘pet.cat.Cat’>”

注：如果此时重新dump类对象，pet.cat.Cat便是随类对象一起保存至文件中的类定义引用。

补充

sys.modules[“cat”] = pet.cat 操作终究了权宜之计，更加建议将之前保存的文件通过此法加载之后，再dump一次，然后及时将sys.modules[“cat”]删除。有如下脚本：

import sys
from types import ModuleType
import joblib
def update_module_path_in_pickled_object(pickle_path: str, old_module_path: str, new_module: ModuleType) -> None:
    """Update a python module's dotted path in a pickle dump if the    corresponding file was renamed.
    Args: pickle_path (str): Path to the pickled object.
    old_module_path (str): The old.dotted.path.to.renamed.module.
    new_module (ModuleType): from new.location import module.
    """
    
    sys.modules[old_module_path] = new_module
    dic = joblib.load(pickle_path)
    del sys.modules[old_module_path]
    joblib.dump(dic, "./test.pkl")

通过上述脚本，便可以将”过期“的二进制文件修改为适应当前文件目录结果的新文件。

补充2

针对本文问题描述中展示的问题，其实还可以通过修改sys.path来load类对象文件，代码如下：

import sys
sys.path.append("./pet")
import joblib
my_cat = joblib.load("my_cat.pkl")
print(type(my_cat))

只是，此时print(type(my_cat))打印的输出为<class ‘cat.Cat’>，因为我直接将模块cat.py的路径放到了sys.path中，而joblib.load(“my_cat.pkl”)背后要import的模块正是cat.py这个模块。

但对于这个问题，修改sys.path并非正道，且还有重大隐患。 为什么并非正道呢？因为sys.path并不能解决此类所有问题。试想，如果我开始时，Cat类定义在animal.cat.py中，调整目录结构之后，改为pet.cat.py中。此时便没法通过简单修改sys.path来实现的。毕竟sys.path修改的只是包搜索的路径，而如果我连之前的包animal都删除了（改为pet了），上哪去搜索呢？要知道，dump的二进制文件可是保存了animal.cat.Cat这样一个完整的路径的。

另外，如果修改sys.path, 这将在不同的包位置创建同一模块的多个副本。这就可能导致更严重的问题，比如，我创建两个Cat对象，这两个Cat对象便可能在会在类型判断的时候不一致，因为这两个类对象的定义可能是来自两个不同的副本。

参考

• pickle - Python pickling after changing a module’s directory - Stack Overflow
• What’s with ML software and pickles? - Made of Bugs (nelhage.com)