深入解析 torchvision 的 create_feature_extractor：用 FX 符号追踪提取模型任意中间层特征

深入解析 torchvision 的 create_feature_extractor用 FX 符号追踪提取模型任意中间层特征【免费下载链接】visionDatasets, Transforms and Models specific to Computer Vision项目地址: https://gitcode.com/gh_mirrors/vi/visiontorchvision.models.feature_extraction是 torchvision 提供的一组模型检视工具它基于 PyTorch 的 FX 符号追踪机制允许你在不改动原始模型代码的前提下接入模型的任意中间层把某个模块或某个算子如relu、add的输出作为特征图取出来。本文以 docs/source/feature_extraction.rst 为骨架结合 feature_extraction.py 源码与 test_backbone_utils.py 测试用例完整讲解节点命名约定、get_graph_node_names与create_feature_extractor的全部参数语义并给出从 ResNet-50 单模型特征提取到 MaskRCNN FPN 骨干改造的端到端可运行代码读完即可在可视化、图像检索与检测骨干定制等场景中直接落地。一、feature_extraction 包能做什么三大典型应用场景torchvision.models.feature_extraction包的核心能力是在不修改原模型前向逻辑的前提下访问输入在模型中的中间变换结果。原文档列举了三类在计算机视觉中非常有代表性的应用可视化特征图Visualizing feature maps把某一卷积层的输出以热力图形式展示辅助理解网络学到了什么提取图像描述子Image descriptors例如为人脸识别、拷贝检测copy-detection、图像检索等任务计算中间层特征向量把选定特征喂给下游子网络做端到端训练例如把一组多尺度特征层级hierarchy of features传给带目标检测头的Feature Pyramid NetworkFPN——这正是检测模型改造骨干网络的经典做法。该包对外只暴露两个公开 API见 feature_extraction.py 的__all__create_feature_extractor真正创建特征提取器get_graph_node_names开发期工具用于查看模型中有哪些节点可供提取。二、工作原理create_feature_extractor 的四步流水线create_feature_extractor的实现机制是符号追踪 计算图重写原文档给出的执行步骤可以精炼为四步符号追踪模型Symbolic tracing用一个 tracer 逐算子执行一遍模型前向得到输入是如何一步步被变换的图表示把用户选定的图节点设为输出将被选中的中间节点挂到output节点上删除冗余节点移除所有位于输出节点下游、不再需要的节点即死代码消除从结果图生成 Python 代码把图打包进一个 PyTorch 模块fx.GraphModule图与代码同时保留后续forward直接执行生成的代码。这一过程本质上是对 PyTorch 官方torch.fx中Tracer、Graph、GraphModule等组件的封装torch.fx提供了更通用、更底层的符号追踪解释此处不再展开外部链接具体实现可对照本仓库源码阅读。值得强调的是torchvision 并没有直接用原生fx.Tracer而是实现了一个NodePathTracer继承自LeafModuleAwareTracer而后者继承自fx.Tracer它在追踪每个操作时会额外把该操作来自模块层级中的哪个位置记录到node_to_qualname字典中——这正是自定义节点命名约定的来源详见下一节。三、节点命名约定如何在模块层级中精确定位中间层要用create_feature_extractor指定输出节点必须先理解本包使用的节点命名规则它与torch.fx的命名略有差异。节点名是一个用.分隔的路径从顶层模块一路走到叶子算子或叶子模块。例如 ResNet-50 中的layer4.2.relu表示ResNet模块第 4 个 stagelayer4里第 2 个 block索引 2中 ReLU 的输出。3.1 截断名truncated node name快捷方式指定返回节点时可以只提供节点名的前缀作为快捷方式。比如用get_graph_node_names打印 ResNet-50 的节点时layer4相关的最后一个节点是layer4.2.relu_2。此时你既可以直接写layer4.2.relu_2也可以只写layer4——按照约定截断名指向该路径下按执行顺序最后一个节点。⚠️ 使用截断名的注意点当一个模块有多个输出时最后一个执行的节点不一定对应你想要的输出。例如 FPN 骨干中某个层可能同时输出用于分类和用于回归的特征。因此建议在依赖截断名之前先阅读输入模型的源码确认最后一个操作确实是你想要的输出。3.2 重复操作的_{int}后缀如果同一个模块或同一个操作在forward中被多次执行节点名会追加_{int}后缀以区分。例如同一个forward方法里出现了三次加法就会得到path.to.module.add、path.to.module.add_1、path.to.module.add_2。关键细节是计数器只在直接父级作用域内维护。所以 ResNet-50 里会同时存在layer4.1.add和layer4.2.add——因为两次加法分别位于不同的 block 中不需要后缀也能区分。3.3 源码级实现后缀计数如何实现父级作用域这一看似简单的约定在 feature_extraction.py 的NodePathTracer._get_node_qualname中有值得注意的实现细节对于call_module之外的算子如add、relu等函数调用在模块限定名之后追加节点本身的名字先去掉 torch.fx 自动加的全局_{int}后缀通过正则._[0-9]$匹配后rsplit(_, 1)再以直接父级为作用域重新编号从已有节点名中反向查找与{node_qualname}或{node_qualname}_{int}匹配的项取最大编号 1编号从 1 开始。也就是说torch.fx 原本的后缀是全局计数的torchvision 通过先剥离、再局部重编号把它改造成了按父级作用域计数的语义。3.4 测试如何锁定这一约定test_backbone_utils.py 的TestSubModule/TestModule结构和一个固定期望列表test_module_nodes精确锁定了该命名约定。TestModule的forward依次执行子模块内两次add、两次relu顶层两次add、两次relu期望的节点名序列为[x, submodule.add, submodule.add_1, submodule.relu, submodule.relu_1, add, add_1, relu, relu_1]test_node_name_conventionstest_backbone_utils.py断言get_graph_node_names的返回与该列表完全一致验证了重复操作加后缀、且计数按父级作用域隔离两条规则。此外test_ops.py 在多处如第 168、867、1178、2474、2566 行对torchvision.ops中的算子调用get_graph_node_names做节点命名一致性检查。四、get_graph_node_names查看可用节点的开发工具get_graph_node_names(model, tracer_kwargsNone, suppress_diff_warningFalse, concrete_argsNone)是一个开发期工具返回(train_nodes, eval_nodes)两个列表分别是对模型在训练模式和评估模式下追踪得到的节点名按执行顺序排列。4.1 为什么节点名不能直接从代码读出来源码 docstringfeature_extraction.py指出了两个根本原因并非所有子模块都会被追踪展开torch.nn下的模块如ReLU、Conv2d属于叶子模块追踪时只产生一个引用该模块forward的节点其内部结构不展开重复操作会被追加_{counter}后缀只有真正执行一遍追踪才能确定每个重复操作到底排第几个。4.2 参数说明参数类型说明modelnn.Module想要打印节点名的模型tracer_kwargsdict可选传给NodePathTracer最终透传给torch.fx.Tracer的关键字参数。默认值为{autowrap_modules: (math, torchvision.ops,), leaf_modules: _get_leaf_modules_for_ops()}即自动包装math模块和所有torchvision.ops并把torchvision.ops中所有nn.Module子类视为叶子模块。注意若用户提供tracer_kwargs默认参数会追加到用户字典中而不是被覆盖suppress_diff_warningbool是否抑制 train/eval 图不一致时的警告默认Falseconcrete_argsdict可选不应被当作 Proxy 的具体参数透传给Tracer.trace其 API 可能不被保证4.3 train/eval 双图差异警告模型会被追踪两次一次model.train()、一次model.eval()。_warn_graph_differencesfeature_extraction.py会比较两组节点若完全一致静默返回若 eval 节点是 train 节点的子序列反之亦然会提示哪个是哪个的子序列否则提示两组节点不同并统一建议选择特征提取节点时你可能需要分别为 train 和 eval 模式指定输出节点。典型触发场景是模型中存在依赖self.training的控制流如 Dropout 分支、训练专用的归一化逻辑。常见的快速使用方式import torchvision from torchvision.models.feature_extraction import get_graph_node_names model torchvision.models.resnet18() train_nodes, eval_nodes get_graph_node_names(model) print(train_nodes)get_graph_node_names在返回前会恢复模型原有的训练状态model.train(is_training)不会因为调用它而改变模型模式。五、create_feature_extractor 完整 API 与参数校验create_feature_extractor(model, return_nodesNone, train_return_nodesNone, eval_return_nodesNone, tracer_kwargsNone, suppress_diff_warningFalse, concrete_argsNone)返回一个fx.GraphModule其forward返回一个OrderedDict键是用户指定的输出名值是选中的中间节点输出。5.1 return_nodes 的两种形态dict形态{节点名: 输出字典的键}。例如{layer1: feat1, layer3: feat2}表示提取layer1与layer3的输出并在返回字典中以feat1、feat2为键list形态等价于{str(i): str(i) for i in n}即节点名直接作为返回字典的键。5.2 train/eval 分离的返回节点train_return_nodes训练模式下的返回节点eval_return_nodes评估模式下的返回节点。当二者都指定时return_nodes不得再指定。这组参数用于处理 train/eval 图节点不一致的场景例如 Dropout 在不同模式下产生不同节点。若未指定return_nodes会被深拷贝为 train 与 eval 两份使用。5.3 参数互斥与存在性校验源码级create_feature_extractor 的实现中有一组严格校验测试用例也一一覆盖test_backbone_utils.pyreturn_nodes与train_return_nodes/eval_return_nodes三者均为None→ 抛ValueErrortrain_return_nodes与eval_return_nodes只指定了其中一个 → 抛ValueError同时指定了return_nodes与train_return_nodes→ 抛ValueError指定了不存在的节点 → 抛ValueError错误信息会提示用get_graph_node_names确认节点名并提示可能要用train_return_nodes/eval_return_nodes分别指定。节点的匹配逻辑是前缀匹配只要存在一个可用节点名满足正则^{query}(\.|$)query即被接受随后按module_qualname逐段与query的.分段比对命中后即把该节点挂到output节点上feature_extraction.py。5.4 默认 tracer 设置让 torchvision.ops 开箱即用_set_default_tracer_kwargsfeature_extraction.py会自动把math与torchvision.ops加入autowrap_modules这些模块内的函数调用会被包装为独立节点把torchvision.ops中所有torch.nn.Module子类通过inspect.getmembers(torchvision.ops)收集作为默认leaf_modules避免追踪进入这些算子的内部实现。这保证了绝大多数官方模型与 ops 组合可以开箱即用地被追踪。六、实战示例一为 ResNet-50 提取四个主干层先看最基础的单模型用法。下面的代码来自原文档并补充了可运行细节import torch from torchvision.models import resnet50 from torchvision.models.feature_extraction import get_graph_node_names, create_feature_extractor # 1) 设计特征提取器之前先打印可用节点 m resnet50() train_nodes, eval_nodes get_graph_node_names(resnet50()) # train_nodes / eval_nodes 分别是模型在 train / eval 模式下 # 按执行顺序排列的全部节点名。本例两者相同 # 若模型包含依赖训练模式的控制流二者可能不同。 # 2) 精确指定各主干层的最后一个节点 return_nodes { # 节点名: 输出字典的键 layer1.2.relu_2: layer1, layer2.3.relu_2: layer2, layer3.5.relu_2: layer3, layer4.2.relu_2: layer4, } # 2) 也可以使用截断名等价于取该路径下最后一个执行节点 return_nodes { layer1: layer1, layer2: layer2, layer3: layer3, layer4: layer4, } # 3) 构建特征提取器 extractor create_feature_extractor(m, return_nodesreturn_nodes) out extractor(torch.rand(1, 3, 224, 224)) # out 是 OrderedDict形如 # { # layer1: layer1 的输出 [1, 64, 56, 56], # layer2: layer2 的输出 [1, 128, 28, 28], # layer3: layer3 的输出 [1, 256, 14, 14], # layer4: layer4 的输出 [1, 512, 7, 7], # }源码 docstring 中给出的最小示例feature_extraction.py也验证了输出形状{layer1: feat1, layer3: feat2}在 224×224 输入下得到feat1为[1, 64, 56, 56]、feat2为[1, 256, 14, 14]。七、实战示例二把 ResNet-50 FPN 封装为 MaskRCNN 骨干特征提取最经典的实战用途之一是把它嵌进一个torch.nn.Module中作为检测模型的骨干。下面的完整示例来自原文档演示如何用create_feature_extractor抽出 ResNet-50 的 4 个主干层再接一个 FPN最后喂给 MaskRCNNimport torch import torch.nn as nn from torchvision.models import resnet50 from torchvision.models.feature_extraction import create_feature_extractor from torchvision.models.detection.mask_rcnn import MaskRCNN from torchvision.models.detection.backbone_utils import LastLevelMaxPool from torchvision.ops.feature_pyramid_network import FeaturePyramidNetwork # MaskRCNN 需要一个挂了 FPN 的 backbone class Resnet50WithFPN(nn.Module): def __init__(self): super(Resnet50WithFPN, self).__init__() # 取一个 resnet50 骨干 m resnet50() # 抽出 4 个主干层 # 注意MaskRCNN 对返回节点名有特定要求这里必须用 # 0, 1, 2, 3 这样的键名映射 self.body create_feature_extractor( m, return_nodes{flayer{k}: str(v) for v, k in enumerate([1, 2, 3, 4])}) # 先跑一次 dry run拿到各层通道数用于构造 FPN inp torch.randn(2, 3, 224, 224) with torch.no_grad(): out self.body(inp) in_channels_list [o.shape[1] for o in out.values()] # 构造 FPN输出通道 256额外加一个最大池化层作为顶层 self.out_channels 256 self.fpn FeaturePyramidNetwork( in_channels_list, out_channelsself.out_channels, extra_blocksLastLevelMaxPool()) def forward(self, x): x self.body(x) x self.fpn(x) return x # 组装完整模型 model MaskRCNN(Resnet50WithFPN(), num_classes91).eval()几个值得注意的实战要点返回字典的键名MaskRCNN的 backbone 期望body(x)返回一个以0、1、2、3为键的字典因此这里用{flayer{k}: str(v) for v, k in enumerate([1, 2, 3, 4])}生成{layer1: 0, layer2: 1, layer3: 2, layer4: 3}映射dry run 求通道数由于不同 ResNet 变体的各层通道数不同先跑一次前向、从输出张量的shape[1]读取in_channels_list再构造FeaturePyramidNetwork是最稳妥的做法LastLevelMaxPool来自 backbone_utils.py为 FPN 增加最后一个尺度的池化分支。若想了解官方更完整的 backbone FPN 组装方式可对照同一文件中的BackboneWithFPN与resnet_fpn_backbone。测试 test_backbone_utils.py 验证了resnet_fpn_backbone的输出键为[0, 1, 2, 3, pool]并校验了trainable_layers与returned_layers的取值范围可作为自定义骨干的参数参考。八、进阶技巧leaf module、autowrap function 与不可追踪代码并非所有模型都能直接 FX 追踪但绝大多数问题可以通过以下手段调教以下三条建议与示例均来自源码 docstringfeature_extraction.py把有问题的子模块变成 leaf module若不需要追踪某个子模块内部把它加入tracer_kwargs的leaf_modules列表追踪器就不会展开它而是让图直接引用该模块的forward把函数变成 leaf function类似地把某些函数加入autowrap_functions让追踪器把整个函数调用当作一个节点而不深入内建 Python 函数通常有坑例如追踪时调用int会直接抛错。可以把它们包进自定义函数再通过autowrap_functions传入。完整示例可直接运行验证import torch import torch.nn as nn from torchvision.models.feature_extraction import create_feature_extractor def leaf_function(x): # 如果被追踪展开int(x) 会抛 TypeError return int(x) class LeafModule(nn.Module): def forward(self, x): # 同样int(x.shape[0]) 若被追踪会抛 TypeError int(x.shape[0]) return torch.nn.functional.relu(x 4) class MyModule(nn.Module): def __init__(self): super().__init__() self.conv nn.Conv2d(3, 1, 3) self.leaf_module LeafModule() def forward(self, x): leaf_function(x.shape[0]) x self.conv(x) return self.leaf_module(x) model create_feature_extractor( MyModule(), return_nodes[leaf_module], tracer_kwargs{leaf_modules: [LeafModule], autowrap_functions: [leaf_function]})测试 test_leaf_module_and_function 对上述模式做了断言生成图中不包含relu节点说明LeafModule未被展开且包含leaf_module节点同时验证了前向与反向out[leaf_module].float().mean().backward()都能正常运行。提示由于默认tracer_kwargs会被追加到用户提供的字典上而非覆盖你在传leaf_modules时不必担心丢失torchvision.ops的默认叶子模块。九、train/eval 双图机制与训练模式切换create_feature_extractor返回的模块是DualGraphModulefeature_extraction.py它同时持有train 与 eval 两张计算图train(modeTrue)时把self.graph切换为train_grapheval()即train(False)时切换为eval_graph只有在模式真正变化时才切换图避免不必要的重复编译子模块get_attr/call_module节点会按两张图合并复制保证两张图都能正常执行保留训练钩子等属性并实现了__deepcopy__非回归测试对应 test_deepcopy修复了 efficientnet 等模型的 deepcopy 问题。为什么需要双图因为有些模型在 train/eval 下计算图不同。测试 test_train_eval 构造了一个极简模型forward中在self.training为真时执行add、为假时执行mul同时dropout(p1.0)的行为也随模式变化。测试分别指定train_return_nodes[dropout, add, sub]与eval_return_nodes[dropout, mul, sub]并断言train 模式下输出中包含add且不含muldropout输出为 0被全部置零eval 模式下输出中包含mul且不含adddropout输出为 1无论从 train 还是 eval 模式开始构造模型与提取器都会保持原始训练状态。与 IntermediateLayerGetter 的等价性仓库中还存在一个更早期的方案IntermediateLayerGetter见 torchvision/models/_utils.py。测试 test_feature_extraction_methods_equivalence 对 ResNet-18 分别用两种方式提取layer1layer4断言两者的参数名与参数值完全一致两者的输出张量逐元素相等。这说明create_feature_extractor是IntermediateLayerGetter的功能超集还支持算子级节点、train/eval 分离节点、leaf module 等是官方推荐的新方案。可脚本化与可序列化测试 test_jit_forward_backward 验证了提取器可以直接torch.jit.script并完成前向/反向test_build_fx_feature_extractor 则用models.list_models(models)遍历了仓库中全部模型随机抽取节点构造提取器验证 list/dict 两种形态、参数互斥校验与错误提示均符合预期。十、在 torchvision.ops 算子上的应用get_graph_node_names不只用于网络模型也常用于检视torchvision.ops中的算子模块。例如 test_ops.py 在多处调用get_graph_node_names(op_obj)如第 168、867、1178、2474、2566 行附近对DeformConv2d、PSROIAlign、PSROIPool、RoIAlign、RoIPool等算子的节点命名进行断言确保算子的节点结构符合预期。这进一步印证了本工具的通用性任何可追踪的nn.Module都可以成为特征提取的对象。总结torchvision.models.feature_extraction以torch.fx符号追踪为基础通过create_feature_extractor与get_graph_node_names两个 API为模型检视提供了标准化的中间层访问方案节点命名.分隔的模块路径 截断名快捷方式 按父级作用域计数的_{int}后缀双图机制DualGraphModule自动维护 train/eval 两张图model.train()/model.eval()无缝切换可追踪性兜底leaf_modules、autowrap_functions、autowrap_modules三项机制解决绝大多数不可追踪代码实战价值特征图可视化、图像检索描述子、检测/分割骨干定制如本文的 MaskRCNN FPN 示例均可直接复用。相关源码与测试均可在当前仓库中继续研读核心实现见 torchvision/models/feature_extraction.py全面的测试覆盖见 test/test_backbone_utils.py算子级应用见 test/test_ops.py旧方案对比对象IntermediateLayerGetter见 torchvision/models/_utils.py。【免费下载链接】visionDatasets, Transforms and Models specific to Computer Vision项目地址: https://gitcode.com/gh_mirrors/vi/vision创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考