FIx traffic splitter example (#1597)

Author: vishalbollu

examples/traffic-splitter/cortex.yaml

@@ -1,20 +1,21 @@
 # WARNING: you are on the master branch; please refer to examples on the branch corresponding to your `cortex version` (e.g. for version 0.22.*, run `git checkout -b 0.22` or switch to the `0.22` branch on GitHub)
 
-- name: iris-classifier-onnx
+- name: iris-classifier-pytorch
   kind: RealtimeAPI
   predictor:
-    type: onnx
-    path: onnx_predictor.py
-    model_path: s3://cortex-examples/onnx/iris-classifier/
+    type: python
+    path: pytorch_predictor.py
+    config:
+      model: s3://cortex-examples/pytorch/iris-classifier/weights.pth
   monitoring:
     model_type: classification
 
-- name: iris-classifier-tf
+- name: iris-classifier-onnx
   kind: RealtimeAPI
   predictor:
-    type: tensorflow
-    path: tensorflow_predictor.py
-    model_path: s3://cortex-examples/tensorflow/iris-classifier/nn/
+    type: onnx
+    path: onnx_predictor.py
+    model_path: s3://cortex-examples/onnx/iris-classifier/
   monitoring:
     model_type: classification
 
@@ -23,5 +24,5 @@
   apis:
     - name: iris-classifier-onnx
       weight: 30
-    - name: iris-classifier-tf
+    - name: iris-classifier-pytorch
       weight: 70

examples/traffic-splitter/model.py

@@ -0,0 +1,59 @@
+# WARNING: you are on the master branch; please refer to examples on the branch corresponding to your `cortex version` (e.g. for version 0.22.*, run `git checkout -b 0.22` or switch to the `0.22` branch on GitHub)
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from sklearn.datasets import load_iris
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score
+
+
+class IrisNet(nn.Module):
+    def __init__(self):
+        super(IrisNet, self).__init__()
+        self.fc1 = nn.Linear(4, 100)
+        self.fc2 = nn.Linear(100, 100)
+        self.fc3 = nn.Linear(100, 3)
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, X):
+        X = F.relu(self.fc1(X))
+        X = self.fc2(X)
+        X = self.fc3(X)
+        X = self.softmax(X)
+        return X
+
+
+if __name__ == "__main__":
+    iris = load_iris()
+    X, y = iris.data, iris.target
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.8, random_state=42)
+
+    train_X = Variable(torch.Tensor(X_train).float())
+    test_X = Variable(torch.Tensor(X_test).float())
+    train_y = Variable(torch.Tensor(y_train).long())
+    test_y = Variable(torch.Tensor(y_test).long())
+
+    model = IrisNet()
+
+    criterion = nn.CrossEntropyLoss()
+
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+    for epoch in range(1000):
+        optimizer.zero_grad()
+        out = model(train_X)
+        loss = criterion(out, train_y)
+        loss.backward()
+        optimizer.step()
+
+        if epoch % 100 == 0:
+            print("number of epoch {} loss {}".format(epoch, loss))
+
+    predict_out = model(test_X)
+    _, predict_y = torch.max(predict_out, 1)
+
+    print("prediction accuracy {}".format(accuracy_score(test_y.data, predict_y.data)))
+
+    torch.save(model.state_dict(), "weights.pth")

examples/traffic-splitter/pytorch_predictor.py

@@ -0,0 +1,50 @@
+# WARNING: you are on the master branch; please refer to examples on the branch corresponding to your `cortex version` (e.g. for version 0.22.*, run `git checkout -b 0.22` or switch to the `0.22` branch on GitHub)
+
+import re
+import torch
+import os
+import boto3
+from botocore import UNSIGNED
+from botocore.client import Config
+from model import IrisNet
+
+labels = ["setosa", "versicolor", "virginica"]
+
+
+class PythonPredictor:
+    def __init__(self, config):
+        # download the model
+        bucket, key = re.match("s3://(.+?)/(.+)", config["model"]).groups()
+
+        if os.environ.get("AWS_ACCESS_KEY_ID"):
+            s3 = boto3.client("s3")  # client will use your credentials if available
+        else:
+            s3 = boto3.client("s3", config=Config(signature_version=UNSIGNED))  # anonymous client
+
+        s3.download_file(bucket, key, "/tmp/model.pth")
+
+        # initialize the model
+        model = IrisNet()
+        model.load_state_dict(torch.load("/tmp/model.pth"))
+        model.eval()
+
+        self.model = model
+
+    def predict(self, payload):
+        # Convert the request to a tensor and pass it into the model
+        input_tensor = torch.FloatTensor(
+            [
+                [
+                    payload["sepal_length"],
+                    payload["sepal_width"],
+                    payload["petal_length"],
+                    payload["petal_width"],
+                ]
+            ]
+        )
+
+        # Run the prediction
+        output = self.model(input_tensor)
+
+        # Translate the model output to the corresponding label string
+        return labels[torch.argmax(output[0])]