模型检查的特征提取¶
该 torchvision.models.feature_extraction
包含特征提取实用程序,使我们能够利用模型来访问输入的中间转换。这可能对计算机视觉中的各种应用很有用。以下是一些示例:
可视化特征图。
提取特征以计算图像描述符,用于人脸识别、复制检测或图像检索等任务。
将选定的特征传递到下游子网络,以便针对特定任务进行端到端训练。例如,将特征层次结构传递到具有目标检测头的特征金字塔网络。
Torchvision 提供 create_feature_extractor()
用于此目的。它的工作原理大致遵循以下步骤:
以符号方式跟踪模型以获取其逐步转换输入的图形表示。
将用户选择的图形节点设置为输出。
删除所有冗余节点(输出节点下游的任何内容)。
从生成的图形中生成 python 代码,并将其与图形本身一起捆绑到 PyTorch 模块中。
该 torch.fx 文档 提供了上述过程和符号跟踪内部工作原理的更通用和详细的解释。
关于节点名称
为了指定哪些节点应该作为提取特征的输出节点,人们应该熟悉此处使用的节点命名约定(这与 torch.fx
中使用的约定略有不同)。节点名称被指定为 .
分隔的路径,从顶层模块向下遍历模块层次结构到叶操作或叶模块。例如 "layer4.2.relu"
在 ResNet-50 中表示 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 提取特征的示例:
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 参考¶
|
创建一个新的图形模块,该模块将给定模型的中间节点作为字典返回,其中用户指定的键为字符串,请求的输出为值。 |
|
开发工具,用于按执行顺序返回节点名称。 |