Skip to content

Commit 9e8afeb

Browse files
vishalbolluvishalbollu
authored
Pytorch example (#477)
1 parent 70c1886 commit 9e8afeb

File tree

6 files changed

+218
-5
lines changed

6 files changed

+218
-5
lines changed

‎examples/image-classifier/alexnet.ipynb‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
{
2+
"nbformat": 4,
3+
"nbformat_minor": 0,
4+
"metadata": {
5+
"colab": {
6+
"name": "alexnet.ipynb",
7+
"provenance": [],
8+
"collapsed_sections": []
9+
},
10+
"kernelspec": {
11+
"name": "python3",
12+
"display_name": "Python 3"
13+
}
14+
},
15+
"cells": [
16+
{
17+
"cell_type": "markdown",
18+
"metadata": {
19+
"id": "_KePiywrVHG2",
20+
"colab_type": "text"
21+
},
22+
"source": [
23+
"# Export Alexnet from Torchvision Models\n",
24+
"In this notebook we convert Alexnet to ONNX and upload it to S3 where it can be deployed by Cortex\n",
25+
"\n",
26+
"Based on: [PytorchOnnxExport](https://github.com/onnx/tutorials/blob/master/tutorials/PytorchOnnxExport.ipynb)"
27+
]
28+
},
29+
{
30+
"cell_type": "markdown",
31+
"metadata": {
32+
"id": "dzphLNy5VswD",
33+
"colab_type": "text"
34+
},
35+
"source": [
36+
"## Install dependencies"
37+
]
38+
},
39+
{
40+
"cell_type": "code",
41+
"metadata": {
42+
"id": "N69aGD72Is4t",
43+
"colab_type": "code",
44+
"colab": {}
45+
},
46+
"source": [
47+
"!pip install torch==1.2.* torchvision==0.4.*"
48+
],
49+
"execution_count": 0,
50+
"outputs": []
51+
},
52+
{
53+
"cell_type": "markdown",
54+
"metadata": {
55+
"id": "2raEvUmojKhK",
56+
"colab_type": "text"
57+
},
58+
"source": [
59+
"## Download and Export Model\n",
60+
"Download the pretrained Alexnet Model and export to ONNX model format:"
61+
]
62+
},
63+
{
64+
"cell_type": "code",
65+
"metadata": {
66+
"id": "zKuFyRTlJUkd",
67+
"colab_type": "code",
68+
"colab": {}
69+
},
70+
"source": [
71+
"import torch\n",
72+
"import torch.onnx\n",
73+
"import torchvision\n",
74+
"\n",
75+
"# Standard ImageNet input - 3 channels, 224x224,\n",
76+
"# values don't matter since we only care about network structure.\n",
77+
"dummy_input = torch.randn(1, 3, 224, 224)\n",
78+
"\n",
79+
"# We are going to use a Pretrained alexnet model\n",
80+
"model = torchvision.models.alexnet(pretrained=True)\n",
81+
"\n",
82+
"# Export to ONNX\n",
83+
"torch.onnx.export(model, dummy_input, \"alexnet.onnx\")"
84+
],
85+
"execution_count": 0,
86+
"outputs": []
87+
},
88+
{
89+
"cell_type": "markdown",
90+
"metadata": {
91+
"id": "4YvEPLmljaMT",
92+
"colab_type": "text"
93+
},
94+
"source": [
95+
"## Upload the model to AWS\n",
96+
"Cortex loads models from AWS, so we need to upload the exported model.\n",
97+
"\n",
98+
"Set these variables to configure your AWS credentials and model upload path:"
99+
]
100+
},
101+
{
102+
"cell_type": "code",
103+
"metadata": {
104+
"id": "y-SAhUH-Jlo_",
105+
"colab_type": "code",
106+
"cellView": "form",
107+
"colab": {}
108+
},
109+
"source": [
110+
"AWS_ACCESS_KEY_ID = \"\" #@param {type:\"string\"}\n",
111+
"AWS_SECRET_ACCESS_KEY = \"\" #@param {type:\"string\"}\n",
112+
"S3_UPLOAD_PATH = \"s3://my-bucket/image-classifier/alexnet.onnx\" #@param {type:\"string\"}\n",
113+
"\n",
114+
"import sys\n",
115+
"import re\n",
116+
"\n",
117+
"if AWS_ACCESS_KEY_ID == \"\":\n",
118+
" print(\"\\033[91m{}\\033[00m\".format(\"ERROR: Please set AWS_ACCESS_KEY_ID\"), file=sys.stderr)\n",
119+
"\n",
120+
"elif AWS_SECRET_ACCESS_KEY == \"\":\n",
121+
" print(\"\\033[91m{}\\033[00m\".format(\"ERROR: Please set AWS_SECRET_ACCESS_KEY\"), file=sys.stderr)\n",
122+
"\n",
123+
"else:\n",
124+
" try:\n",
125+
" bucket = re.search(\"s3://(.+?)/\", S3_UPLOAD_PATH).group(1)\n",
126+
" key = re.search(\"s3://.+?/(.+)\", S3_UPLOAD_PATH).group(1)\n",
127+
" except:\n",
128+
" print(\"\\033[91m{}\\033[00m\".format(\"ERROR: Invalid s3 path (should be of the form s3://my-bucket/path/to/file)\"), file=sys.stderr)"
129+
],
130+
"execution_count": 0,
131+
"outputs": []
132+
},
133+
{
134+
"cell_type": "markdown",
135+
"metadata": {
136+
"id": "HmvoV7v96jip",
137+
"colab_type": "text"
138+
},
139+
"source": [
140+
"Upload the model to S3:"
141+
]
142+
},
143+
{
144+
"cell_type": "code",
145+
"metadata": {
146+
"id": "--va3L2KNBHX",
147+
"colab_type": "code",
148+
"colab": {}
149+
},
150+
"source": [
151+
"import boto3\n",
152+
"\n",
153+
"s3 = boto3.client(\"s3\", aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY)\n",
154+
"print(\"Uploading {} ...\".format(S3_UPLOAD_PATH), end = '')\n",
155+
"s3.upload_file(\"alexnet.onnx\", bucket, key)\n",
156+
"print(\" ✓\")"
157+
],
158+
"execution_count": 0,
159+
"outputs": []
160+
},
161+
{
162+
"cell_type": "markdown",
163+
"metadata": {
164+
"id": "acHZMDxqjnNQ",
165+
"colab_type": "text"
166+
},
167+
"source": [
168+
"<!-- CORTEX_VERSION_MINOR -->\n",
169+
"That's it! See the [example on GitHub](https://github.com/cortexlabs/cortex/tree/master/examples/alexnet) for how to deploy the model as an API."
170+
]
171+
}
172+
]
173+
}

‎examples/image-classifier/alexnet_handler.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
import requests
2+
import numpy as np
3+
import base64
4+
from PIL import Image
5+
from io import BytesIO
6+
from torchvision import transforms
7+
8+
labels = requests.get(
9+
"https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt"
10+
).text.split("\n")[1:]
11+
12+
13+
# https://github.com/pytorch/examples/blob/447974f6337543d4de6b888e244a964d3c9b71f6/imagenet/main.py#L198-L199
14+
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
15+
preprocess = transforms.Compose(
16+
[transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize]
17+
)
18+
19+
20+
def pre_inference(sample, metadata):
21+
if "url" in sample:
22+
image = requests.get(sample["url"]).content
23+
elif "base64" in sample:
24+
image = base64.b64decode(sample["base64"])
25+
26+
img_pil = Image.open(BytesIO(image))
27+
img_tensor = preprocess(img_pil)
28+
img_tensor.unsqueeze_(0)
29+
return img_tensor.numpy()
30+
31+
32+
def post_inference(prediction, metadata):
33+
return labels[np.argmax(np.array(prediction).squeeze())]

‎examples/image-classifier/cortex.yaml‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,15 @@
22
name: image
33

44
- kind: api
5-
name: classifier
6-
model: s3://cortex-examples/inception
7-
request_handler: handler.py
5+
name: classifier-inception
6+
model: s3://cortex-examples/image-classifier/inception
7+
request_handler: inception_handler.py
8+
tracker:
9+
model_type: classification
10+
11+
- kind: api
12+
name: classifier-alexnet
13+
model: s3://cortex-examples/image-classifier/alexnet.onnx
14+
request_handler: alexnet_handler.py
815
tracker:
916
model_type: classification

‎examples/image-classifier/inception.ipynb‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,7 @@
136136
"source": [
137137
"AWS_ACCESS_KEY_ID = \"\" #@param {type:\"string\"}\n",
138138
"AWS_SECRET_ACCESS_KEY = \"\" #@param {type:\"string\"}\n",
139-
"S3_UPLOAD_PATH = \"s3://my-bucket/inception\" #@param {type:\"string\"}\n",
139+
"S3_UPLOAD_PATH = \"s3://my-bucket/image-classifier/inception\" #@param {type:\"string\"}\n",
140140
"\n",
141141
"import sys\n",
142142
"import re\n",

‎examples/image-classifier/handler.py‎ renamed to ‎examples/image-classifier/inception_handler.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
from PIL import Image
55
from io import BytesIO
66

7-
87
labels = requests.get(
98
"https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt"
109
).text.split("\n")

‎examples/image-classifier/requirements.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
pillow
2+
torchvision

0 commit comments

Comments
 (0)