# 模型检查的特征提取

> 原文：[`pytorch.org/vision/stable/feature_extraction.html`](https://pytorch.org/vision/stable/feature_extraction.html)

`torchvision.models.feature_extraction`包含特征提取工具，让我们可以访问模型对输入的中间转换。这对计算机视觉中的各种应用可能很有用。只是一些例子包括：

+   可视化特征图。

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

+   将选定的特征传递给下游子网络，以便根据特定任务进行端到端训练。例如，将特征的层次结构传递给具有目标检测头的特征金字塔网络。

Torchvision 提供了`create_feature_extractor()`用于此目的。它大致按照以下步骤进行：

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

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

1.  删除所有冗余节点（输出节点之后的所有节点）。

1.  从生成的图形中生成 Python 代码，并将其与图形一起捆绑到 PyTorch 模块中。

[torch.fx 文档](https://pytorch.org/docs/stable/fx.html)提供了对上述过程和符号跟踪内部工作的更一般和详细的解释。

**关于节点名称**

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

+   在为`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 提取特征的示例：

```py
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 参考

| `create_feature_extractor`(model[, ...]) | 创建一个新的图模块，将给定模型的中间节点作为字典返回，用户可以指定键作为字符串，请求的输出作为值。 |
| --- | --- |
| `get_graph_node_names`(model[, tracer_kwargs, ...]) | 开发工具，按执行顺序返回节点名称。 |