initial commit

Author: kykosic

.dockerignore

@@ -0,0 +1 @@
+**/target

.gitignore

@@ -0,0 +1,16 @@
+# MacOS cache
+.DS_Store
+
+# Python cache
+__pycache__
+
+# Editor settings
+.vscode
+.ipynb_checkpoints
+
+# Rust compilation targets
+server/target
+
+# Downloaded data
+data/
+

Dockerfile

@@ -0,0 +1,28 @@
+FROM rust:latest AS build
+
+WORKDIR /app
+COPY server/Cargo.toml server/Cargo.lock ./
+COPY server/src/ src/
+RUN cargo build --release
+RUN find . -type d | grep "libtorch/lib$" | xargs -I{} mv {} libtorch
+
+FROM debian:buster-slim
+
+RUN apt-get update \
+    && apt-get install -y \
+        ca-certificates \
+        libgomp1 \
+    && apt-get autoremove \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+COPY --from=build /app/target/release/actix-torch-server /usr/local/bin/actix-torch-server
+copy --from=build /app/libtorch/* /usr/lib/
+COPY saved_model/ /app/saved_model/
+
+RUN useradd -mU -s /bin/bash actix
+USER actix
+
+EXPOSE 8080
+ENTRYPOINT ["actix-torch-server", "--model-path=/app/saved_model/model.pt"]

README.md

@@ -0,0 +1,72 @@
+# Serving PyTorch with Actix-Web
+
+This repository gives an example of training a machine learning model using PyTorch, exporting that model, then serving inference over a RESTful API using Actix-Web in rust.
+
+Implemented by request based on the original [TensorFlow example](https://github.com/kykosic/actix-tensorflow-example).
+
+
+For more information on the tools used, check out the following repositories:
+* [PyTorch](https://pytorch.org/)
+* [Torch Rust (tch-rs)](https://github.com/LaurentMazare/tch-rs)
+* [Actix Web](https://github.com/actix/actix-web)
+
+
+## Overview
+The repository has 3 sections:
+* `./training` – Contains the script which trains a neural network to recognize digits.
+* `./server` – Contains the RESTful API webserver rust code.
+* `./client` – Contains a sample script that demonstrates making a request to the server.
+
+The training script will output a saved neural network model in compiled TorchScript format. The server then loads this into memory on startup. The server accepts a JSON payload at the `/mnist` endpoint with a single key "image" that is a base64 encoded image (PNG or JPG). This image is decoded, rescaled to the correct input dimensions, converted to grayscale, normalized (matching the training data normalization), and finally submitted to the model for inference. Predictions are returned with a "label" integer value and a "confidence" float between 0 and 1.
+
+## Setup
+This example assumes you have [rust installed](https://www.rust-lang.org/tools/install) and python 3.6+ setup. To install the needed python dependencies:
+```
+pip install -r requirements.txt
+```
+
+## Training
+The model used is a simple convolutional neural network trained on the [MNIST dataset](http://yann.lecun.com/exdb/mnist/). The data is automatically downloaded using the `torchvision` library. To train the model:
+```
+python training/train.py
+```
+This will output a saved model to `./saved_model/model.pt`. A pre-trained model is included in this repository. The model isn't too large and can be trained without any GPU.
+
+## Serving
+The server code is a rust crate located in `./server`. In order to run, the server requires the saved model directory location specified with the `--model-path` command line argument. you can try running the server with:
+```
+cd server
+cargo run -- --model-dir ../saved_model/model.pt
+```
+
+## Serving in Docker
+For actual deployments, you probably would want to build a release in a container to serve the API. To build the docker image:
+```
+docker build -t actix-torch .
+```
+Then to run the image locally for testing:
+```
+docker run --rm -it -p 8080:8080 actix-torch
+```
+
+## Client Testing
+With the server running locally, you can test inference using `./client/client.py`. Included is a PNG file with a handwritten "3" that is base64-encoded and submitted to the server.
+
+To test:
+```
+python client/client.py
+```
+
+Input:
+
+<img src="client/test_image_3.png" width="80">
+
+Expected output:
+```
+POST to http://127.0.0.1:8080/mnist
+Response (200)
+Content: {
+    "label": 3,
+    "confidence": 0.9999999
+}
+```

client/client.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python
+"""
+    Example python client for testing the prediction server
+"""
+import base64
+import json
+import os
+
+import requests
+
+
+if __name__ == '__main__':
+    url = "http://127.0.0.1:8080/mnist"
+    this_dir = os.path.dirname(os.path.abspath(__file__))
+    image_path = os.path.join(this_dir, 'test_image_3.png')
+
+    print(f"Reading image from {image_path}")
+    with open(image_path, 'rb') as f:
+        image = f.read()
+
+    data = {
+        'image': base64.b64encode(image).decode()
+    }
+    print(f"POST to {url}")
+    res = requests.post(url,
+                        json=data,
+                        headers={'content-type': 'application/json'})
+    print(f"Response ({res.status_code})")
+    if res.status_code == 200:
+        print(f"Content: {json.dumps(res.json(), indent=4)}")
+    else:
+        print(f"Body: {res.text}")

client/test_image_3.png

@@ -0,0 +1,3 @@
+requests~=2.26.0
+torch~=1.9.0
+torchvision~=0.10.0