SSD
模型描述
此 SSD300 模型基於《SSD: Single Shot MultiBox Detector》論文，該論文將 SSD 描述為“一種使用單個深度神經網路檢測影像中物體的方法”。輸入大小固定為 300×300。
此模型與論文中描述的模型的主要區別在於主幹網路。具體來說，VGG 模型已過時，並被 ResNet-50 模型取代。
根據《Speed/accuracy trade-offs for modern convolutional object detectors》論文，對主幹網路進行了以下增強：
- 原始分類模型中的 conv5_x、avgpool、fc 和 softmax 層已被移除。
- conv4_x 中的所有步幅都設定為 1×1。
主幹網路之後是 5 個額外的卷積層。除了卷積層之外，我們還附加了 6 個檢測頭。
- 第一個檢測頭連線到最後一個 conv4_x 層。
- 其他五個檢測頭連線到相應的 5 個額外層。
檢測頭與論文中引用的類似，但在每次卷積之後都透過額外的 BatchNorm 層進行了增強。
示例
在下面的示例中，我們將使用預訓練的 SSD 模型來檢測樣本影像中的物體並可視化結果。
要執行此示例，您需要安裝一些額外的 Python 包。這些包用於影像預處理和視覺化。
pip install numpy scipy scikit-image matplotlib
載入在 COCO 資料集上預訓練的 SSD 模型，以及一組用於方便和全面格式化模型輸入和輸出的實用方法。
import torch
ssd_model = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd')
utils = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_ssd_processing_utils')
現在，準備載入的模型進行推理。
ssd_model.to('cuda')
ssd_model.eval()
準備用於物體檢測的輸入影像。（下面的示例連結對應於 COCO 資料集中的前幾張測試影像，但您也可以在此處指定本地影像的路徑）
uris = [
'http://images.cocodataset.org/val2017/000000397133.jpg',
'http://images.cocodataset.org/val2017/000000037777.jpg',
'http://images.cocodataset.org/val2017/000000252219.jpg'
]
格式化影像以符合網路輸入並將其轉換為張量。
inputs = [utils.prepare_input(uri) for uri in uris]
tensor = utils.prepare_tensor(inputs)
執行 SSD 網路以執行物體檢測。
with torch.no_grad():
detections_batch = ssd_model(tensor)
預設情況下，每個輸入影像的 SSD 網路的原始輸出包含 8732 個帶有定位和類機率分佈的框。讓我們過濾此輸出，以更全面的格式只獲取合理的檢測（置信度>40%）。
results_per_input = utils.decode_results(detections_batch)
best_results_per_input = [utils.pick_best(results, 0.40) for results in results_per_input]
該模型在 COCO 資料集上訓練，我們需要訪問該資料集才能將類 ID 轉換為物體名稱。首次下載註釋可能需要一些時間。
classes_to_labels = utils.get_coco_object_dictionary()
最後，讓我們視覺化我們的檢測結果。
from matplotlib import pyplot as plt
import matplotlib.patches as patches
for image_idx in range(len(best_results_per_input)):
fig, ax = plt.subplots(1)
# Show original, denormalized image...
image = inputs[image_idx] / 2 + 0.5
ax.imshow(image)
# ...with detections
bboxes, classes, confidences = best_results_per_input[image_idx]
for idx in range(len(bboxes)):
left, bot, right, top = bboxes[idx]
x, y, w, h = [val * 300 for val in [left, bot, right - left, top - bot]]
rect = patches.Rectangle((x, y), w, h, linewidth=1, edgecolor='r', facecolor='none')
ax.add_patch(rect)
ax.text(x, y, "{} {:.0f}%".format(classes_to_labels[classes[idx] - 1], confidences[idx]*100), bbox=dict(facecolor='white', alpha=0.5))
plt.show()
詳情
有關模型輸入和輸出、訓練配方、推理和效能的詳細資訊，請訪問：github 和/或 NGC
參考文獻
