torchvision.models.feature_extraction

特征提取工具让我们能够访问模型对输入进行的中间转换。这在计算机视觉的多种应用中都非常有用。例如：

• 可视化特征图。

• 提取特征以计算图像描述符，用于面部识别、复制检测或图像检索等任务。

• 将选定的特征传递给下游子网络，进行端到端训练，针对特定任务。例如，将特征层次传递给具有目标检测头的特征金字塔网络。

Torchvision为此提供了create_feature_extractor()。 它的工作流程大致如下：

1. 通过符号追踪模型，逐步获取其如何转换输入的图形表示。

2. 将用户选择的图节点设置为输出。

3. 移除所有冗余节点（输出节点下游的任何内容）。

4. 从生成的图生成 Python 代码，并将其与图本身一起打包到一个 PyTorch 模块中。

The torch.fx文档 提供了上述过程和符号追踪内部工作原理的更通用和详细的解释。

关于节点名称

为了指定哪些节点应作为提取特征的输出节点，您需要熟悉此处使用的节点命名约定（与torch.fx中略有不同）。节点名称被指定为一个用.分隔的路径，从顶级模块向下遍历模块层次结构，直到叶子操作或叶子模块。例如，在ResNet-50中，"layer4.2.relu"表示第4层第2个块中的ReLU的输出。ResNet模块。以下是一些需要注意的细节：

• 在指定节点名称时，可以为create_feature_extractor()提供一个节点名称的简写版本作为快捷方式。要了解这如何工作，请尝试创建一个ResNet-50模型，并使用train_nodes, _ = get_graph_node_names(model) print(train_nodes)打印节点名称，观察与layer4相关的最后一个节点是"layer4.2.relu_2"。可以指定"layer4.2.relu_2"作为返回节点，或者只需指定"layer4"，因为按照惯例，这指的是layer4执行顺序中的最后一个节点。

• 如果某个模块或操作被重复使用超过一次，节点名称会添加一个额外的 _{int} 后缀以区分。例如，在同一个 forward 方法中，加法操作 (+) 可能被使用了三次。那么就会有 "path.to.module.add"、"path.to.module.add_1"、"path.to.module.add_2"。计数器在直接父级的作用域内维护。因此，在 ResNet-50 中有一个 "layer4.1.add" 和一个 "layer4.2.add"。由于加法操作位于不同的块中，不需要后缀来区分。

一个示例

这是一个我们如何为MaskRCNN提取特征的示例：

import torch
from torchvision.models import resnet50
from torchvision.models.feature_extraction import get_graph_node_names
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


# To assist you in designing the feature extractor you may want to print out
# the available nodes for resnet50.
m = resnet50()
train_nodes, eval_nodes = get_graph_node_names(resnet50())

# The lists returned, are the names of all the graph nodes (in order of
# execution) for the input model traced in train mode and in eval mode
# respectively. You'll find that `train_nodes` and `eval_nodes` are the same
# for this example. But if the model contains control flow that's dependent
# on the training mode, they may be different.

# To specify the nodes you want to extract, you could select the final node
# that appears in each of the main layers:
return_nodes = {
    # node_name: user-specified key for output dict
    'layer1.2.relu_2': 'layer1',
    'layer2.3.relu_2': 'layer2',
    'layer3.5.relu_2': 'layer3',
    'layer4.2.relu_2': 'layer4',
}

# But `create_feature_extractor` can also accept truncated node specifications
# like "layer1", as it will just pick the last node that's a descendent of
# of the specification. (Tip: be careful with this, especially when a layer
# has multiple outputs. It's not always guaranteed that the last operation
# performed is the one that corresponds to the output you desire. You should
# consult the source code for the input model to confirm.)
return_nodes = {
    'layer1': 'layer1',
    'layer2': 'layer2',
    'layer3': 'layer3',
    'layer4': 'layer4',
}

# Now you can build the feature extractor. This returns a module whose forward
# method returns a dictionary like:
# {
#     'layer1': output of layer 1,
#     'layer2': output of layer 2,
#     'layer3': output of layer 3,
#     'layer4': output of layer 4,
# }
create_feature_extractor(m, return_nodes=return_nodes)

# Let's put all that together to wrap resnet50 with MaskRCNN

# MaskRCNN requires a backbone with an attached FPN
class Resnet50WithFPN(torch.nn.Module):
    def __init__(self):
        super(Resnet50WithFPN, self).__init__()
        # Get a resnet50 backbone
        m = resnet50()
        # Extract 4 main layers (note: MaskRCNN needs this particular name
        # mapping for return nodes)
        self.body = create_feature_extractor(
            m, return_nodes={f'layer{k}': str(v)
                             for v, k in enumerate([1, 2, 3, 4])})
        # Dry run to get number of channels for 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()]
        # Build FPN
        self.out_channels = 256
        self.fpn = FeaturePyramidNetwork(
            in_channels_list, out_channels=self.out_channels,
            extra_blocks=LastLevelMaxPool())

    def forward(self, x):
        x = self.body(x)
        x = self.fpn(x)
        return x


# Now we can build our model!
model = MaskRCNN(Resnet50WithFPN(), num_classes=91).eval()

API 参考

torchvision.models.feature_extraction.create_feature_extractor(model: torch.nn.modules.module.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: Dict = {}, suppress_diff_warning: bool = False) → torch.fx.graph_module.GraphModule

创建一个新的图模块，该模块从给定模型返回中间节点，并以用户指定的字符串键作为字典的键，请求的输出作为值。通过 FX 重写模型的计算图，使其返回所需的节点作为输出。所有未使用的节点及其对应的参数都将被移除。

指定的输出节点必须作为从顶层模块到叶操作或叶模块的.分隔路径指定。有关此处使用的节点命名约定的更多详细信息，请参见相关小标题中的文档。

并非所有模型都可以进行 FX 迹迹化，尽管通过一些调整可以使它们配合工作。这里有一份（不详尽）的技巧清单：

• If you don’t need to trace through a particular, problematic sub-module, turn it into a “leaf module” by passing a list of leaf_modules as one of the tracer_kwargs (see example below). It will not be traced through, but rather, the resulting graph will hold a reference to that module’s forward method.

• Likewise, you may turn functions into leaf functions by passing a list of autowrap_functions as one of the tracer_kwargs (see example below).

• Some inbuilt Python functions can be problematic. For instance, int will raise an error during tracing. You may wrap them in your own function and then pass that in autowrap_functions as one of the tracer_kwargs.

有关FX的更多信息，请参阅 torch.fx文档。

Parameters

• 模型 (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）。

• suppress_diff_warning (bool, 可选) – 是否在训练版本和评估版本的图之间存在差异时抑制警告。默认为False。

Examples:

>>> # 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]})

torchvision.models.feature_extraction.get_graph_node_names(model: torch.nn.modules.module.Module, tracer_kwargs: Dict = {}, suppress_diff_warning: bool = False) → Tuple[List[str], List[str]]

开发工具，按执行顺序返回节点名称。有关节点名称的说明，请参见create_feature_extractor()下的注释。这对于查看哪些节点名称可用于特征提取非常有用。模型的代码中节点名称难以直接读取的原因有两个：

1. Not all submodules are traced through. Modules from torch.nn all fall within this category.

2. Nodes representing the repeated application of the same operation or leaf module get a _{counter} postfix.

该模型被追踪了两次：一次在训练模式下，另一次在评估模式下。两种节点名称都会被返回。

有关此处使用的节点命名约定的更多详细信息，请参见 相关小标题 在 文档 中。

Parameters

• 模型 (nn.Module) – 我们希望打印节点名称的模型

• tracer_kwargs (字典, 可选) –

  一个关键字参数字典用于 NodePathTracer （它们最终会被传递给 torch.fx.Tracer）。

• suppress_diff_warning (bool, 可选) – 是否在训练版本和评估版本的图之间存在差异时抑制警告。默认为False。

Returns

从以训练模式跟踪模型得到的一组节点名称，以及从以评估模式跟踪模型得到的另一组节点名称。

Return type

元组(列表, 列表)

Examples:

>>> model = torchvision.models.resnet18()
>>> train_nodes, eval_nodes = get_graph_node_names(model)