如何在不使用dill且无自定义类直接引用时用torch.save/load存取自定义对象

核心报错逻辑 #

torch.save 默认基于 Python 标准库 pickle 实现序列化，这个问题和你的类是不是神经网络模型无关，只要是被torch.save序列化的Python自定义类，都遵循pickle的通用规则：它不会把自定义类的代码本身存入序列化文件，仅会在文件中记录类的全路径引用（格式为所在模块.类名）。执行torch.load时，反序列化器会按照这个路径，在当前运行环境的全局命名空间中查找对应的类定义，找到后才会基于类定义、配合存储的实例属性值重建对象。

你遇到的报错本质是：当初执行torch.save存储DagDataPreparation对象时，这个类的定义位于当时的主运行脚本（也就是__main__模块）中，序列化器记录的类路径为__main__.DagDataPreparation。当你切换到另一个脚本执行加载操作时，当前的__main__模块已经是新的运行脚本，里面不存在DagDataPreparation的定义，反序列化器找不到目标属性，就抛出了AttributeError。

两个能正常运行的方案的底层原理 #

• 传入dill作为序列化后端：dill是pickle的扩展实现，会直接把类的代码块完整序列化存入文件，加载时直接从文件中读取类定义重建对象，完全不需要依赖当前环境的类查找逻辑，因此不会出现找不到类的问题。
• 显式导入类：执行from data_lib.dag.dataset_preparation import DagDataPreparation后，Python会把这个类注册到当前运行环境的全局命名空间中，反序列化器遍历全局属性时能匹配到对应的类对象，就可以正常完成加载。

跨文件加载的报错普遍性 #

如果序列化对象时类是定义在主运行脚本下（即模块路径为__main__），所有未提前将该类导入到当前运行环境全局命名空间的开发者，执行加载操作时都会遇到这个报错。

你之前没有遇到过同类问题，核心原因是过去的存储、加载操作都是在同一个脚本内完成的：存储时类在__main模块下，加载时同一个__main模块依然保留着类定义，环境完全匹配自然不会报错。一旦将存储、加载逻辑拆到两个独立脚本中，这个问题就会立刻暴露。

不依赖dill、无需手动导入类的解决方案 #

• 从根源修正类的存放位置：所有需要被torch序列化的自定义类，必须定义在独立的、可被正常导入的模块文件中，绝对不能直接写在作为程序入口的主运行脚本里。只要存储对象时，类的引用是来自独立模块（如data_lib.dag.dataset_preparation.DagDataPreparation）而非__main__，序列化文件中记录的类路径就是可跨环境访问的标准模块路径，后续加载时只要项目的模块搜索路径配置正确，反序列化器会自动顺着模块路径找到类定义，完全不需要手动导入。

  注意：即使你把类移到了独立模块里，如果存储时直接在主脚本中重新定义了同名类、没有从独立模块导入，序列化器依然会把类路径记为__main__.DagDataPreparation，还是会触发同类报错。

• 自定义反序列化器的类查找映射：如果已经生成了旧的序列化文件（类路径被记为__main__.DagDataPreparation），不想重新生成文件，也不想在每个加载脚本里手动导入类，可以自定义一个继承自torch.serialization.Unpickler的类，重写find_class方法：当反序列化器查找__main模块下的DagDataPreparation时，自动将其映射到独立模块中存放的实际类定义，实现无手动导入的加载逻辑。
• 统一全局模块映射：在项目的加载逻辑入口处，将独立模块中的类自动挂载到__main__的命名空间下，让反序列化器查找时可以直接匹配到对应类，不需要在每个单独的加载脚本中重复编写导入语句。

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