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-0 | GPUNet-0 在 GV100 上具有最快的測量延遲 |
GPUNet-1 | GPUNet-1 在 GPUNet-0 的基礎上增加了一層,提高了準確性 |
GPUNet-2 | GPUNet-2 在 GPUNet-0 的基礎上增加了兩層,準確性更高 |
GPUNet-P0 | GPUNet-P0 是蒸餾模型,準確性高於 GPUNet-0 但延遲相似 |
GPUNet-P1 | GPUNet-P1 是蒸餾模型,準確性甚至高於 GPUNet-1 但延遲相似 |
GPUNet-D1 | GPUNet-D1 在所有 GPUNet 中準確性排名第二 |
GPUNet-D2 | GPUNet-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
參考