使用 PyTorch-Neuron 和 AWS Neuron 编译器 - 深度学习 AMI
AWS 文档中描述的 AWS 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅中国的 AWS 服务入门

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

使用 PyTorch-Neuron 和 AWS Neuron 编译器

-Neuron 编译 API 提供了一种方法,用于编译您可在 PyTorch 设备上运行的模型图。AWS Inferentia

经过训练的模型必须先编译为 Inferentia 目标,才能部署在 Inf1 实例上。以下教程将编译 torchvision ResNet50 模型并将其导出为已保存的 TorchScript 模块。此模型随后将用于运行推理。

为方便起见,本教程使用 Inf1 实例进行编译和推理。在实际操作中,您也可以使用其他实例类型来编译模型,例如 c5 实例系列。然后,您必须将已编译的模型部署到 Inf1 推理服务器中。有关更多信息,请参阅 AWS Neuron PyTorch 开发工具包文档

Prerequisites

使用本教程之前,您应已完成 将 DLAMI 与 AWS Neuron 结合使用 中的设置步骤。您还应该熟悉深度学习知识以及如何使用 DLAMI。

激活 Conda 环境

使用以下命令激活 PyTorch-Neuron Conda 环境:

source activate aws_neuron_pytorch_p36

要退出当前 Conda 环境,请运行:

source deactivate

Resnet50 编译

创建一个名为 pytorch_trace_resnet50.py 的 Python 脚本,并在其中包含以下内容。此脚本使用 PyTorch-Neuron 编译 Python API 来编译 ResNet-50 模型。

import torch import numpy as np import os import torch_neuron from torchvision import models image = torch.zeros([1, 3, 224, 224], dtype=torch.float32) ## Load a pretrained ResNet50 model model = models.resnet50(pretrained=True) ## Tell the model we are using it for evaluation (not training) model.eval() model_neuron = torch.neuron.trace(model, example_inputs=[image]) ## Export to saved model model_neuron.save("resnet50_neuron.pt")

运行编译脚本。

python pytorch_trace_resnet50.py

编译将需要几分钟的时间。编译完成后,已编译的模型将在本地目录中另存为 resnet50_neuron.pt

ResNet50 推理

创建一个名为 pytorch_infer_resnet50.py 的 Python 脚本,并在其中包含以下内容。此脚本会下载一个示例映像,然后使用该映像对已编译的模型运行推理过程。

import os import time import torch import torch_neuron import json import numpy as np from urllib import request from torchvision import models, transforms, datasets ## Create an image directory containing a small kitten os.makedirs("./torch_neuron_test/images", exist_ok=True) request.urlretrieve("https://raw.githubusercontent.com/awslabs/mxnet-model-server/master/docs/images/kitten_small.jpg", "./torch_neuron_test/images/kitten_small.jpg") ## Fetch labels to output the top classifications request.urlretrieve("https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json","imagenet_class_index.json") idx2label = [] with open("imagenet_class_index.json", "r") as read_file: class_idx = json.load(read_file) idx2label = [class_idx[str(k)][1] for k in range(len(class_idx))] ## Import a sample image and normalize it into a tensor normalize = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) eval_dataset = datasets.ImageFolder( os.path.dirname("./torch_neuron_test/"), transforms.Compose([ transforms.Resize([224, 224]), transforms.ToTensor(), normalize, ]) ) image, _ = eval_dataset[0] image = torch.tensor(image.numpy()[np.newaxis, ...]) ## Load model model_neuron = torch.jit.load( 'resnet50_neuron.pt' ) ## Predict results = model_neuron( image ) # Get the top 5 results top5_idx = results[0].sort()[1][-5:] # Lookup and print the top 5 labels top5_labels = [idx2label[idx] for idx in top5_idx] print("Top 5 labels:\n {}".format(top5_labels) )

使用以下命令对已编译模型运行推理过程:

python pytorch_infer_resnet50.py

您的输出应与以下内容类似:

Top 5 labels: ['tiger', 'lynx', 'tiger_cat', 'Egyptian_cat', 'tabby']