模型中心

GPUNet

模型描述

GPUNets 是 NVIDIA 推出的全新系列卷积神经网络，专为部署和生产环境打造，旨在最大化 NVIDIA GPU 和 TensorRT 的性能。

GPUNet 由 NVIDIA AI 团队使用创新的神经架构搜索（NAS）方法构建，展现了顶尖的推理性能，速度比 EfficientNet-X 和 FBNet-V3 快达 2 倍。本笔记本允许您加载并测试我们在 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