模型描述

GPUNets 是英伟达推出的一系列新型、可部署且生产就绪的卷积神经网络，通过自动化设计最大化发挥英伟达 GPU 和 TensorRT 的性能。

GPUNet 由英伟达 AI 采用新颖的神经架构搜索(NAS)方法精心打造，推理性能达到最先进水平，比 EfficientNet-X 和 FBNet-V3 快达 2 倍。本 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 数据集上预训练的英伟达 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) helper 函数根据模型挑选 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

参考资料