模型中心

GPUNet

模型描述

GPUNets 是 NVIDIA 推出的一系列新型部署和生产就绪的卷积神经网络，它们经过自动设计，旨在最大限度地发挥 NVIDIA GPU 和 TensorRT 的性能。

GPUNet 由 NVIDIA AI 采用新颖的神经架构搜索 (NAS) 方法精心打造，展示了最先进的推理性能，比 EfficientNet-X 和 FBNet-V3 快两倍。本 Notebook 允许您加载和测试我们CVPR-2022 论文中列出的所有 GPUNet 模型实现。您可以使用此 Notebook 快速加载列出的每个模型以执行推理运行。

示例

在下面的示例中，默认加载预训练的 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