使用create_feature_extractor时遇TraceError: Proxy无法迭代问题求助

问题场景 #

使用create_feature_extractor提取Faster RCNN模型中间特征时，触发以下错误：

TraceError: Proxy object cannot be iterated. This can be attempted when the Proxy is used in a loop or as a *args or **kwargs function argument. See the torch.fx docs on pytorch.org for a more detailed explanation of what types of control flow can be traced, and check out the Proxy docstring for help troubleshooting Proxy iteration errors

对应代码：

from torchvision.models.feature_extraction import create_feature_extractor
from torchvision.models.detection import fasterrcnn_mobilenet_v3_large_320_fpn

m = fasterrcnn_mobilenet_v3_large_320_fpn()

return_nodes = {
    'layer1': 'layer1'
}

create_feature_extractor(m, return_nodes=return_nodes)
return_nodes

错误原因 #

create_feature_extractor依赖torch.fx的静态追踪机制，但Faster RCNN这类检测模型内部包含动态控制流（如基于检测结果的循环、条件判断），这类逻辑无法被fx静态解析，导致生成的Proxy对象在迭代时触发错误。

解决方案 #

方案1：单独对Backbone使用create_feature_extractor #

检测模型的backbone（特征提取主干）是纯静态卷积结构，无动态控制流，可单独提取后使用create_feature_extractor：

from torchvision.models.feature_extraction import create_feature_extractor
from torchvision.models.detection import fasterrcnn_mobilenet_v3_large_320_fpn

# 初始化检测模型
m = fasterrcnn_mobilenet_v3_large_320_fpn()
# 分离出backbone模块
backbone = m.backbone

return_nodes = {
    'layer1': 'layer1'
}

# 对backbone创建特征提取器
feature_extractor = create_feature_extractor(backbone, return_nodes=return_nodes)

方案2：使用前向钩子（Hook）提取特征 #

如果需要保留完整检测模型流程并提取中间特征，可通过手动注册前向钩子绕过fx追踪限制：

from torchvision.models.detection import fasterrcnn_mobilenet_v3_large_320_fpn
import torch

m = fasterrcnn_mobilenet_v3_large_320_fpn()
m.eval()

# 用于存储提取到的特征
extracted_features = {}

# 定义钩子函数，捕获目标层的输出
def capture_feature(module, input, output):
    extracted_features['layer1'] = output

# 定位到目标层并注册前向钩子
target_layer = m.backbone.layer1
hook_handle = target_layer.register_forward_hook(capture_feature)

# 输入符合模型要求的示例张量
dummy_input = [torch.randn(3, 320, 320)]
m(dummy_input)

# 查看提取到的特征形状
print(extracted_features['layer1'].shape)

# 使用完毕后移除钩子，避免内存泄漏
hook_handle.remove()

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