模型中心

EfficientNet

模型描述

EfficientNet 是一个图像分类模型系列，首次在《EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks》中提出。本 Notebook 允许您加载并测试 EfficientNet-B0、EfficientNet-B4、EfficientNet-WideSE-B0 和 EfficientNet-WideSE-B4 模型。

EfficientNet-WideSE 模型使用的 Squeeze-and-Excitation (SE) 层比原始的 EfficientNet 模型更宽，SE 模块的宽度与深度可分离卷积（Depthwise Separable Convolutions）的宽度成比例，而不是与块（block）宽度成比例。

WideSE 模型的准确率略高于原始模型。

该模型在 Volta 和 NVIDIA Ampere GPU 架构上使用 Tensor Core 进行混合精度训练。因此，研究人员获得的训练速度比不使用 Tensor Core 的训练快 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