模型中心

EfficientNet

模型描述

EfficientNet 是一個影像分類模型家族。它首次在EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks中描述。本筆記允許您載入和測試 EfficientNet-B0、EfficientNet-B4、EfficientNet-WideSE-B0 和 EfficientNet-WideSE-B4 模型。

EfficientNet-WideSE 模型使用的 Squeeze-and-Excitation 層比原始 EfficientNet 模型更寬，SE 模組的寬度與 Depthwise Separable Convolutions 的寬度成正比，而不是與塊寬度成正比。

WideSE 模型比原始模型略微準確。

此模型在 Volta 和 NVIDIA Ampere GPU 架構上使用 Tensor Cores 進行混合精度訓練。因此，研究人員可以獲得比不使用 Tensor Cores 訓練快 2 倍以上的結果，同時體驗混合精度訓練的好處。此模型針對每個 NGC 每月容器釋出進行測試，以確保長期保持一致的準確性和效能。

在使用混合精度進行訓練時，我們使用 NHWC 資料佈局。

示例

在下面的示例中，我們將使用預訓練的 EfficientNet 模型對影像執行推理並呈現結果。

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

!pip install validators matplotlib

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

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
print(f'Using {device} for inference')

載入在 ImageNet 資料集上預訓練的模型。

您可以選擇以下模型

TorchHub 入口點描述
nvidia_efficientnet_b0基線 EfficientNet
nvidia_efficientnet_b4縮放後的 EfficientNet
nvidia_efficientnet_widese_b0Squeeze-and-Excitation 層比基線 EfficientNet 模型更寬的模型
nvidia_efficientnet_widese_b4Squeeze-and-Excitation 層比縮放後的 EfficientNet 模型更寬的模型

還有模型的量化版本，但它們需要 nvidia 容器。請參閱量化模型

efficientnet = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_efficientnet_b0', pretrained=True)
utils = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_convnets_processing_utils')

efficientnet.eval().to(device)

準備示例輸入資料。

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)

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

with torch.no_grad():
    output = torch.nn.functional.softmax(efficientnet(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 和/或 NGC

參考資料

EfficientNets 是一系列影像分類模型，它們實現了最先進的精度，同時尺寸和速度提高了一個數量級。使用 Tensor Cores 進行混合精度訓練。

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