使用 PyTorch-Neuron 和 AWS Neuron 編譯器 - AWS 深度學習 AMIs

本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。

使用 PyTorch-Neuron 和 AWS Neuron 編譯器

PyTorch-Neuron 編譯API提供一種方法來編譯模型圖表,您可以在 AWS Inferentia 裝置上執行。

經過訓練的模型必須編譯至 Inferentia 目標,才能部署到 Inf1 執行個體上。下列教學課程會編譯 torchvision ResNet50 模型,並將其匯出為儲存的 TorchScript 模組。接著,即可使用此模型執行推論。

為了方便起見,本教學課程使用 Inf1 執行個體進行編譯和推論。實際上,您可以使用 c5 執行個體系列等其他執行個體類型來編譯模型。接著,您必須將已編譯的模型部署到 Inf1 推論伺服器。如需詳細資訊,請參閱 AWS Neuron PyTorch SDK 文件

必要條件

在使用本教學課程之前,您應該已完成 使用 AWS Neuron 啟動DLAMI執行個體 中的設置步驟。您也應該熟悉深度學習和使用 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 模型。

注意

編撰 torchvision 模型時,您應該注意的 torchvision 版本與 torch 套件之間存在相依性。這些相依性規則可以透過 pip 管理。Torchvision==0.6.1 符合 torch==1.5.1 版本,而 torchvision==0.8.2 符合 torch==1.7.1 版本。

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']