Python PyTorch keypointrcnn_resnet50_fpn用法及代碼示例

本文簡要介紹python語言中 torchvision.models.detection.keypointrcnn_resnet50_fpn 的用法。

用法:

torchvision.models.detection.keypointrcnn_resnet50_fpn(pretrained=False, progress=True, num_classes=2, num_keypoints=17, pretrained_backbone=True, trainable_backbone_layers=None, **kwargs)

參數：

• pretrained(bool) -如果為 True，則返回在 COCO train2017 上預訓練的模型

• progress(bool) -如果為 True，則顯示下載到 stderr 的進度條

• num_classes(int) -模型的輸出類數(包括背景)

• num_keypoints(int) -關鍵點數，默認 17

• pretrained_backbone(bool) -如果為 True，則返回一個在 Imagenet 上預訓練過主幹的模型

• trainable_backbone_layers(int) -從最終塊開始的可訓練(未凍結)resnet 層數。有效值介於 0 和 5 之間，其中 5 表示所有主幹層都是可訓練的。

使用 ResNet-50-FPN 主幹構造 Keypoint R-CNN 模型。

參考：“Mask R-CNN”。

模型的輸入應該是一個張量列表，每個形狀為 [C, H, W] ，每個圖像一個，並且應該在 0-1 範圍內。不同的圖像可以有不同的尺寸。

模型的行為取決於它是處於訓練模式還是評估模式。

在訓練期間，模型需要輸入張量以及目標(字典列表)，其中包含：

• 框 ( FloatTensor[N, 4] )： [x1, y1, x2, y2] 格式的 ground-truth 框，包含 0 <= x1 < x2 <= W 和 0 <= y1 < y2 <= H 。

• labels (Int64Tensor[N])：每個ground-truth框的類標簽

• 關鍵點 ( FloatTensor[N, K, 3] )：每個 N 實例的 K 關鍵點位置，格式為 [x, y, visibility] ，其中 visibility=0 表示關鍵點不可見。

該模型在訓練期間返回 Dict[Tensor]，包含 RPN 和 R-CNN 的分類和回歸損失，以及關鍵點損失。

在推理過程中，模型隻需要輸入張量，並將後處理的預測作為 List[Dict[Tensor]] 返回，每個輸入圖像一個。 Dict的字段如下，其中N是檢測到的實例數：

• 框 ( FloatTensor[N, 4] )：[x1, y1, x2, y2] 格式的預測框，包含 0 <= x1 < x2 <= W 和 0 <= y1 < y2 <= H 。

• labels (Int64Tensor[N])：每個實例的預測標簽

• 分數 (Tensor[N])：每個實例的分數

• 關鍵點 ( FloatTensor[N, K, 3] )：預測關鍵點的位置，采用 [x, y, v] 格式。

有關輸出的更多詳細信息，您可以參考實例分割模型。

關鍵點 R-CNN 可導出到 ONNX 以用於固定批量大小，輸入圖像大小固定。

例子：

>>> model = torchvision.models.detection.keypointrcnn_resnet50_fpn(pretrained=True)
>>> model.eval()
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> predictions = model(x)
>>>
>>> # optionally, if you want to export the model to ONNX:
>>> torch.onnx.export(model, x, "keypoint_rcnn.onnx", opset_version = 11)

使用 keypointrcnn_resnet50_fpn 的示例：

可視化實用程序