當前位置: 首頁>>代碼示例 >>用法及示例精選 >>正文


Python PyTorch retinanet_resnet50_fpn用法及代碼示例


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

用法:

torchvision.models.detection.retinanet_resnet50_fpn(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True, trainable_backbone_layers=None, **kwargs)

參數

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

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

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

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

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

使用 ResNet-50-FPN 主幹構建 RetinaNet 模型。

參考:“Focal Loss for Dense Object Detection”

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

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

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

  • 框 ( FloatTensor[N, 4] ): [x1, y1, x2, y2] 格式的 ground-truth 框,包含 0 <= x1 < x2 <= W0 <= y1 < y2 <= H

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

該模型在訓練期間返回 Dict[Tensor],包含分類和回歸損失。

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

  • 框 ( FloatTensor[N, 4] ):[x1, y1, x2, y2] 格式的預測框,包含 0 <= x1 < x2 <= W0 <= y1 < y2 <= H

  • labels (Int64Tensor[N]):每次檢測的預測標簽

  • 分數(Tensor[N]):每次檢測的分數

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

例子:

>>> model = torchvision.models.detection.retinanet_resnet50_fpn(pretrained=True)
>>> model.eval()
>>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
>>> predictions = model(x)

使用 retinanet_resnet50_fpn 的示例:

相關用法


注:本文由純淨天空篩選整理自pytorch.org大神的英文原創作品 torchvision.models.detection.retinanet_resnet50_fpn。非經特殊聲明,原始代碼版權歸原作者所有,本譯文未經允許或授權,請勿轉載或複製。