Uso de PyTorch-Neuron y el compilador de AWS Neuron

La API de compilación de PyTorch-Neuron proporciona un método para compilar un gráfico de modelo que se puede ejecutar en un dispositivo AWS de Inferentia.

Un modelo entrenado debe compilarse en un destino de Inferentia antes de poder implementarlo en instancias Inf1. El siguiente tutorial compila el modelo ResNet50 de torchvision y lo exporta como módulo TorchScript guardado. A continuación, el modelo se utiliza para ejecutar la inferencia.

Para mayor comodidad, el tutorial utiliza una instancia Inf1 tanto para la compilación como para la inferencia. En la práctica, puede compilar el modelo con otro tipo de instancia, como la familia de instancias c5. A continuación, debe implementar el modelo compilado en el servidor de inferencia Inf1. Para obtener más información, consulte la documentación del SDK de AWS Neuron PyTorch.

Requisitos previos

Antes de utilizar este tutorial, debería haber completado los pasos de configuración de Lanzamiento de una instancia de DLAMI con AWS Neuron. También debe estar familiarizado con el aprendizaje profundo y con el uso de la DLAMI.

Activación del entorno Conda

Active el entorno Conda de PyTorch-Neuron mediante el siguiente comando:

source activate aws_neuron_pytorch_p36

Para salir del entorno Conda actual, ejecute:

source deactivate

Compilación de Resnet50

Cree un script de Python denominado pytorch_trace_resnet50.py con el siguiente contenido. Este script usa la API de Python de compilación de PyTorch-Neuron para compilar un modelo 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")

Ejecute el script de compilación.

python pytorch_trace_resnet50.py

La compilación tardará unos minutos. Cuando haya finalizado, el modelo compilado se guardará como resnet50_neuron.pt en el directorio local.

Inferencia ResNet50

Cree un script de Python denominado pytorch_infer_resnet50.py con el siguiente contenido. Este script descarga una imagen de muestra y la utiliza para ejecutar la inferencia con el modelo compilado.

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) )

Ejecute la inferencia con el modelo compilado mediante el siguiente comando:

python pytorch_infer_resnet50.py

El resultado debería tener el siguiente aspecto:

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