create_feature_extractor¶
- torchvision.models.feature_extraction.create_feature_extractor(model: Module, return_nodes: Optional[Union[List[str], Dict[str, str]]] = None, train_return_nodes: Optional[Union[List[str], Dict[str, str]]] = None, eval_return_nodes: Optional[Union[List[str], Dict[str, str]]] = None, tracer_kwargs: Optional[Dict[str, Any]] = None, suppress_diff_warning: bool = False, concrete_args: Optional[Dict[str, Any]] = None) GraphModule[源代码]¶
创建一个新的图形模块，该模块将给定模型中的中间节点作为字典返回，其中用户指定的键为字符串，请求的输出为值。这是通过通过 FX 重写模型的计算图来实现的，以便将所需的节点作为输出返回。所有未使用的节点及其对应的参数都将被删除。
所需的输出节点必须指定为以
.分隔的路径，从顶级模块向下遍历模块层次结构到叶子操作或叶子模块。有关此处使用的节点命名约定的更多详细信息，请参阅相关小标题中的文档。并非所有模型都可被 FX 追踪，尽管经过一些修改，它们可以被使其配合工作。以下是一些（非详尽）提示
如果您不需要遍历特定有问题的子模块，请通过将
leaf_modules列表作为tracer_kwargs之一传递（请参见下面的示例）将其转换为“叶子模块”。它不会被遍历，而是生成的图形将保存对该模块的 forward 方法的引用。同样，您可以通过将
autowrap_functions的列表作为tracer_kwargs之一传递来将函数转换为叶子函数（请参见下面的示例）。某些内置的 Python 函数可能存在问题。例如，
int在跟踪期间会引发错误。您可以将它们包装在您自己的函数中，然后将其作为autowrap_functions之一传递给tracer_kwargs。
有关 FX 的更多信息，请参阅 torch.fx 文档。
- 参数：
model (nn.Module) – 我们将在其上提取特征的模型
return_nodes (列表 或 字典，可选) – 一个
List或一个Dict，包含将返回其激活的节点的名称（或部分名称 - 请参见上文说明）。如果它是Dict，则键是节点名称，值是用户为图模块返回的字典指定的键。如果它是List，则将其视为一个Dict，将节点规范字符串直接映射到输出名称。如果指定了train_return_nodes和eval_return_nodes，则不应指定此参数。train_return_nodes (列表 或 字典，可选) – 类似于
return_nodes。如果训练模式的返回节点与评估模式的返回节点不同，则可以使用此参数。如果指定了此参数，则还必须指定eval_return_nodes，并且不应指定return_nodes。eval_return_nodes (列表 或 字典，可选) – 类似于
return_nodes。如果训练模式的返回节点与评估模式的返回节点不同，则可以使用此参数。如果指定了此参数，则还必须指定train_return_nodes，并且不应指定 return_nodes。tracer_kwargs (字典，可选) –
NodePathTracer的关键字参数字典（将其传递给其父类 torch.fx.Tracer）。默认情况下，它将设置为包装并使所有 torchvision 操作成为叶子节点：{“autowrap_modules”: (math, torchvision.ops,),”leaf_modules”: _get_leaf_modules_for_ops(),} 警告：如果用户提供了 tracer_kwargs，则上述默认参数将附加到用户提供的字典中。suppress_diff_warning (布尔值，可选) – 是否抑制图的训练版本和评估版本之间存在差异时的警告。默认为 False。
concrete_args (Optional[Dict[字符串, 任意]]) – 不应被视为代理的具体参数。根据 Pytorch 文档，此参数的 API 可能无法保证。
示例
>>> # Feature extraction with resnet >>> model = torchvision.models.resnet18() >>> # extract layer1 and layer3, giving as names `feat1` and feat2` >>> model = create_feature_extractor( >>> model, {'layer1': 'feat1', 'layer3': 'feat2'}) >>> out = model(torch.rand(1, 3, 224, 224)) >>> print([(k, v.shape) for k, v in out.items()]) >>> [('feat1', torch.Size([1, 64, 56, 56])), >>> ('feat2', torch.Size([1, 256, 14, 14]))] >>> # Specifying leaf modules and leaf functions >>> def leaf_function(x): >>> # This would raise a TypeError if traced through >>> return int(x) >>> >>> class LeafModule(torch.nn.Module): >>> def forward(self, x): >>> # This would raise a TypeError if traced through >>> int(x.shape[0]) >>> return torch.nn.functional.relu(x + 4) >>> >>> class MyModule(torch.nn.Module): >>> def __init__(self): >>> super().__init__() >>> self.conv = torch.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]})
