模型描述

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

此模型使用 Tensor Cores 在 Volta、Turing 和 NVIDIA Ampere GPU 架构上进行混合精度训练。因此，研究人员可以比不使用 Tensor Cores 的训练速度快 3 倍，同时享受混合精度训练的好处。此模型针对每个 NGC 每月容器发布进行测试，以确保长时间内精度和性能的一致性。

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

模型架构

SEArch

图片来源：Squeeze-and-Excitation Networks

图片显示了 SE 块的架构以及它在 ResNet 瓶颈块中的位置。

请注意，SE-ResNeXt101-32x4d 模型可以使用 TorchScript、ONNX Runtime 或 TensorRT 作为执行后端部署到 NVIDIA Triton Inference Server 上进行推理。详情请查看 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。

参考资料