SE-ResNeXt101

模型描述

SE-ResNeXt101-32x4d 是一款在 ResNeXt101-32x4d 模型基础上,增加了 Squeeze-and-Excitation Networks 论文中所介绍的“挤压与激励”(Squeeze-and-Excitation)模块的模型。

该模型使用混合精度训练,并利用了 Volta、Turing 和 NVIDIA Ampere GPU 架构上的 Tensor Core。因此,研究人员获得结果的速度比不使用 Tensor Core 进行训练快 3 倍,同时还能享受到混合精度训练带来的优势。该模型会针对每个 NGC 每月容器版本进行测试,以确保其准确性和性能随时间的推移保持一致。

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

模型架构

图片来源:Squeeze-and-Excitation Networks

图片展示了 SE 模块的架构以及它在 ResNet 瓶颈块(bottleneck block)中的放置位置。

请注意,SE-ResNeXt101-32x4d 模型可通过 NVIDIA Triton 推理服务器进行推理部署,并使用 TorchScript、ONNX Runtime 或 TensorRT 作为执行后端。详情请查阅 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