fastai加载pkl模型报找不到Tf属性的AttributeError如何解决

你在Jupyter Notebook中训练模型时，Tf函数定义在Jupyter的运行时主命名空间（即__main__模块）下，fastai导出pkl模型时会将该函数的引用按__main__.Tf的路径写入序列化文件。反序列化时pickle只会到当前进程的__main__模块（也就是你启动脚本时的入口文件main.py对应模块）下查找名为Tf的属性，你之前尝试的放置位置都不满足这个查找规则：

• 放在main()函数、if __name__ == '__main__'代码块内：导入train_model时就会触发顶层的load_learner执行，此时上述代码块还未运行，Tf还未被定义
• 放在其他模块、或跨模块导入：仅会将Tf绑定到对应模块的命名空间，不会注册到__main__模块下，pickle无法识别

方案1（最推荐）：调整main.py代码结构，在导入train_model前顶层定义Tf #

必须将Tf定义在main.py的顶层位置，且放在导入train_model的代码之前，保证导入触发模型加载时，__main__模块下已经存在Tf属性。
修改后的main.py代码：

import glob
from pathlib import Path

# 【关键】Tf必须定义在导入train_model之前，且为顶层函数
def Tf(o):
    return '/mnt/scratch2/DLinTHDP/PathLAKE/Version_4_fastai/Dataset/CD8/Train/masks/'+f'{o.stem}_P{o.suffix}'

# 再导入train模块，此时导入触发load_learner可正常找到Tf
from train_model import train

ROOT = Path("folder/path")  # Detection Folder

def main(root: Path):
    train(root)

if __name__ == '__main__':
    main(ROOT)

同时建议调整train_model.py的代码，不要在模块顶层执行模型加载（导入即加载的逻辑灵活性极差，很容易触发顺序问题），把加载逻辑移到函数内部：

from pathlib import Path
from fastai.vision.all import *

folder_path = Path('.')

def train(root: Path):
    # 模型加载移到函数内部，此时主模块的Tf已完成注册
    learn = load_learner('model.pkl', cpu=True)
    learn.load('model_3C_34_CELW_V_1.1')  # weights
    # 后续推理、训练逻辑

方案2：手动注入Tf到__main__模块（无需调整加载位置） #

如果你一定要保留train_model.py顶层加载模型的逻辑，可以在train_model.py中load_learner执行前，手动将Tf绑定到__main__模块：

import sys
from pathlib import Path
from fastai.vision.all import *

folder_path = Path('.')

# 先定义和训练时完全一致的Tf函数
def Tf(o):
    return '/mnt/scratch2/DLinTHDP/PathLAKE/Version_4_fastai/Dataset/CD8/Train/masks/'+f'{o.stem}_P{o.suffix}'

# 手动将Tf注册到主模块命名空间
setattr(sys.modules['__main__'], 'Tf', Tf)

# 再执行模型加载
learn = load_learner('model.pkl', cpu=True)
learn.load('model_3C_34_CELW_V_1.1')  # weights

def train(root: Path):
    # 后续逻辑

注意：Tf的函数名、参数列表必须和训练时完全一致，即使函数内的硬编码路径你当前推理阶段用不到，也不能修改，pickle是按函数的引用身份查找，不是按运行逻辑匹配。

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