GPUNet

模型描述

GPUNets 是 NVIDIA 推出的一系列可部署、生產就緒的卷積神經網路新家族,它們經過自動設計,旨在最大限度地發揮 NVIDIA GPU 和 TensorRT 的效能。

GPUNet 由 NVIDIA AI 採用新穎的神經架構搜尋 (NAS) 方法精心打造,展示了最先進的推理效能,比 EfficientNet-X 和 FBNet-V3 快兩倍。本筆記本允許您載入並測試我們CVPR-2022 論文中列出的所有 GPUNet 模型實現。您可以使用本筆記本快速載入列出的每個模型以執行推理執行。

示例

在下面的示例中,預設載入預訓練的 GPUNet-0 模型,以對影像執行推理並呈現結果。您可以將預設預訓練模型的載入從 GPUNet-0 切換到以下模型之一:GPUNet-1、GPUNet-2、GPUNet-P0、GPUNet-P1、GPUNet-D1 或 GPUNet-D2。

安裝先決條件

要執行此示例,您需要安裝一些額外的 Python 包。這些包用於影像預處理和視覺化。

!pip install validators matplotlib
!pip install timm==0.5.4
import torch
from PIL import Image
import torchvision.transforms as transforms
import numpy as np
import json
import requests
import matplotlib.pyplot as plt
import warnings

warnings.filterwarnings('ignore')
%matplotlib inline


if torch.cuda.is_available():
    device = torch.device("cuda") 
    !nvidia-smi
else:
    device = torch.device("cpu")

print(f'Using {device} for inference')

載入預訓練模型

預設載入在 ImageNet 資料集上預訓練的 NVIDIA GPUNet-0 模型。您可以將預設預訓練模型的載入從 GPUNet-0 切換到下面列出的以下模型之一。

模型架構在載入模型的輸出中可見。有關詳細的架構和延遲資訊,請分別參考原始倉庫中的架構部分和 CVPR-2022 論文中的表3

請選擇以下預訓練模型之一

TorchHub 模型描述
GPUNet-0GPUNet-0 在 GV100 上具有最快的測量延遲
GPUNet-1GPUNet-1 在 GPUNet-0 的基礎上增加了一層,提高了準確性
GPUNet-2GPUNet-2 在 GPUNet-0 的基礎上增加了兩層,準確性更高
GPUNet-P0GPUNet-P0 是蒸餾模型,準確性高於 GPUNet-0 但延遲相似
GPUNet-P1GPUNet-P1 是蒸餾模型,準確性甚至高於 GPUNet-1 但延遲相似
GPUNet-D1GPUNet-D1 在所有 GPUNet 中準確性排名第二
GPUNet-D2GPUNet-D2 在所有 GPUNet 中準確性最高
model_type = "GPUNet-0" # select one from above
precision = "fp32" # select either fp32 of fp16 (for better performance on GPU)

gpunet = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_gpunet', pretrained=True, model_type=model_type, model_math=precision)
utils = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_convnets_processing_utils')

gpunet.to(device)
gpunet.eval()

準備推理資料

為推理準備樣本輸入資料。

uris = [
    'http://images.cocodataset.org/test-stuff2017/000000024309.jpg',
    'http://images.cocodataset.org/test-stuff2017/000000028117.jpg',
    'http://images.cocodataset.org/test-stuff2017/000000006149.jpg',
    'http://images.cocodataset.org/test-stuff2017/000000004954.jpg',
]

batch = torch.cat(
    [utils.prepare_input_from_uri(uri) for uri in uris]
).to(device)

if precision == "fp16":
    batch = batch.half()
    
print("Ready to run inference...")

執行推理

使用 pick_n_best(predictions=output, n=topN) 輔助函式根據模型選擇 N 個最有可能的假設。

with torch.no_grad():
    output = torch.nn.functional.softmax(gpunet(batch), dim=1)
    
results = utils.pick_n_best(predictions=output, n=5)

顯示結果

for uri, result in zip(uris, results):
    img = Image.open(requests.get(uri, stream=True).raw)
    img.thumbnail((256,256), Image.ANTIALIAS)
    plt.imshow(img)
    plt.show()
    print(result)

詳情

有關模型輸入和輸出、訓練方法、推理和效能的詳細資訊,請訪問:github

參考

GPUNet 是一系列新的卷積神經網路,旨在最大化 NVIDIA GPU 和 TensorRT 的效能。

模型型別: 視覺
提交者: NVIDIA