模型检查的特征提取

模型检查的特征提取#

torchvision.models.feature_extraction 包包含了特征提取工具，这些工具能够访问模型的中间变换层，从而获取输入数据的中间特征。这在计算机视觉的各种应用中非常有用。例如：

可视化特征图。
提取特征以计算图像描述符，用于人脸识别、复制检测或图像检索等任务。
将选定的特征传递给下游子网络，进行端到端的特定任务训练。例如，将层次化的特征传递给带有目标检测头的特征金字塔网络。

Torchvision 为此提供了 create_feature_extractor()。其工作原理大致如下：

符号追踪模型，逐步生成计算图表示，展示如何变换输入。
将用户选择的计算图节点设置为输出。
移除所有冗余节点（即输出节点之后的所有节点）。
从生成的计算图中生成 Python 代码，并将其与计算图本身一同打包成 PyTorch 模块。

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

关于节点名称#

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

在指定 feature_extraction() 的节点名称时，您可以提供截断版本的节点名称作为快捷方式。要了解这一点如何工作，请尝试创建 ResNet-50 模型，并使用train_nodes, _ = get_graph_node_names(model) print(train_nodes) 打印节点名称，您会发现与 layer4 相关的最后一个节点是 "layer4.2.relu_2"。您可以将 "layer4.2.relu_2" 指定为返回节点，或者只是 "layer4"，因为这按照惯例指的是 layer4 的执行顺序中的最后一个节点。
如果某个模块或操作重复多次，节点名称会附加一个额外的 _{int} 后缀以消除歧义。例如，也许在同一个前向方法中使用了三次加法（+）运算。那么会有 "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


# 为了帮助你设计特征提取器，你可能想要打印出 resnet50 的可用节点。
m = resnet50()
train_nodes, eval_nodes = get_graph_node_names(resnet50())

# 返回的列表是输入模型在训练模式和评估模式下跟踪的图节点的名称（按执行顺序排列）。
# 你会发现对于这个例子，`train_nodes` 和 `eval_nodes` 是相同的。但如果模型包含依赖于训练模式的控制流，它们可能会有所不同。

# 要指定你想要提取的节点，你可以选择每个主要层中出现的最后一个节点：
return_nodes = {
    # node_name: 用户指定的输出字典键
    'layer1.2.relu_2': 'layer1',
    'layer2.3.relu_2': 'layer2',
    'layer3.5.relu_2': 'layer3',
    'layer4.2.relu_2': 'layer4',
}

# 但 `create_feature_extractor` 也可以接受截断的节点规范，如 "layer1"，因为它会选择规范的最后一个后代节点。
# （提示：使用时要小心，特别是当一个层有多个输出时。不能保证最后一个操作是与你期望的输出相对应的操作。你应该查阅输入模型的源代码以确认。）
return_nodes = {
    'layer1': 'layer1',
    'layer2': 'layer2',
    'layer3': 'layer3',
    'layer4': 'layer4',
}

# 现在你可以构建特征提取器。这将返回一个模块，其前向方法返回一个字典，如下所示：
# {
#     'layer1': 第 1 层的输出,
#     'layer2': 第 2 层的输出,
#     'layer3': 第 3 层的输出,
#     'layer4': 第 4 层的输出,
# }
create_feature_extractor(m, return_nodes=return_nodes)

# 让我们将所有这些结合起来，用 MaskRCNN 包装 resnet50

# MaskRCNN 需要一个带有附加 FPN 的主干网络
class Resnet50WithFPN(torch.nn.Module):
    def __init__(self):
        super(Resnet50WithFPN, self).__init__()
        # 获取一个 resnet50 主干网络
        m = resnet50()
        # 提取 4 个主要层（注意：MaskRCNN 需要这个特定的名称映射用于返回节点）
        self.body = create_feature_extractor(
            m, return_nodes={f'layer{k}': str(v)
                             for v, k in enumerate([1, 2, 3, 4])})
        # 试运行以获取 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
        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


# 现在可以构建模型了！
model = MaskRCNN(Resnet50WithFPN(), num_classes=91).eval()

模型检查的特征提取

目录

模型检查的特征提取#

关于节点名称#

示例#