SE-ResNeXt101

模型描述

SE-ResNeXt101-32x4d 是一个ResNeXt101-32x4d模型,其中添加了在Squeeze-and-Excitation Networks论文中介绍的Squeeze-and-Excitation模块。

该模型在Volta、Turing和NVIDIA Ampere GPU架构上使用Tensor Cores进行混合精度训练。因此,研究人员可以获得比不使用Tensor Cores训练快3倍的结果,同时体验混合精度训练的优势。该模型针对每个NGC月度容器版本进行测试,以确保随着时间的推移保持一致的准确性和性能。

在使用混合精度进行训练时,我们使用 NHWC 数据布局

模型架构

图片来源:Squeeze-and-Excitation Networks

图片展示了SE模块的架构及其在ResNet瓶颈模块中的位置。

请注意,SE-ResNeXt101-32x4d模型可以使用TorchScript、ONNX Runtime或TensorRT作为执行后端,部署到NVIDIA Triton 推理服务器进行推理。有关详细信息,请查看NGC

示例

在下面的示例中,我们将使用预训练的SE-ResNeXt101-32x4d模型对图像执行推理并呈现结果。

要运行此示例,您需要安装一些额外的 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 数据集上预训练的模型。

resneXt = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_se_resnext101_32x4d')
utils = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_convnets_processing_utils')

resneXt.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(resneXt(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

参考文献

添加了 Squeeze-and-Excitation 模块的 ResNeXt,使用 Tensor Cores 进行混合精度训练。

模型类型: 视觉
提交者: NVIDIA