Skip to content

Commit

Permalink
added neuron resource (#39)
Browse files Browse the repository at this point in the history
  • Loading branch information
freschri committed Mar 14, 2024
1 parent d6a68f9 commit 45934f1
Showing 1 changed file with 96 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
apiVersion: v1
kind: ConfigMap
metadata:
name: pytorch-resnet50-script
data:
pytorch-resnet50.sh: |
#!/usr/bin/env python
import torch
import numpy as np
import os
import torch_neuron
from torchvision import models, transforms, datasets
import time
from urllib import request
import json
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])
## 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, ...])
## Predict
while True:
results = model_neuron( image )
time.sleep(0.1)
---
apiVersion: v1
kind: Pod
metadata:
name: pytorch-inference-resnet50
spec:
tolerations:
- key: node-role.kubernetes.io/control-plane
operator: Exists
effect: NoSchedule
- key: node-role.kubernetes.io/master
operator: Exists
effect: NoSchedule
containers:
- name: pytorch
image: 763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference-neuron:1.13.1-neuron-py310-sdk2.13.2-ubuntu20.04
command: ["/script/pytorch-resnet50.sh"]
volumeMounts:
- name: script
mountPath: "/script"
resources:
limits:
cpu: 4
memory: 4Gi
aws.amazon.com/neuron: 1
requests:
cpu: "1"
memory: 1Gi
securityContext:
capabilities:
add:
- IPC_LOCK
volumes:
- name: script
configMap:
name: pytorch-resnet50-script
defaultMode: 0777

0 comments on commit 45934f1

Please sign in to comment.