diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..411837e --- /dev/null +++ b/Dockerfile @@ -0,0 +1,17 @@ +FROM python:3.8-slim-buster + +WORKDIR /docker_demo + +RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple + +RUN pip install torch==1.7.1 torchvision==0.8.2 --extra-index-url https://download.pytorch.org/whl/cpu + +RUN pip install Flask==2.0.3 + +ADD . . + +RUN mkdir -p /root/.cache/torch/hub/checkpoints/ + +RUN mv /docker_demo/densenet121-a639ec97.pth /root/.cache/torch/hub/checkpoints/densenet121-a639ec97.pth + +CMD ["python", "app.py"] \ No newline at end of file diff --git a/README.md b/README.md index cd5fd67..7341487 100644 --- a/README.md +++ b/README.md @@ -4,24 +4,37 @@ This repo contains a sample code to show how to create a Flask API server by dep If you'd like to learn how to deploy to Heroku, then check [this repo](https://github.com/avinassh/pytorch-flask-api-heroku). +## Deploy with Docker -## How to +Build image -Install the dependencies: + docker build --network=host -t pytorch_flask:v0 . - pip install -r requirements.txt +Run Docker + docker run --name pytorch_flask -p 5000:5000 -d pytorch_flask:v0 -Run the Flask server: +send the image file in a request: - FLASK_ENV=development FLASK_APP=app.py flask run + curl -X POST -F file=@cat_pic.jpeg http://localhost:5000/predict -From another tab, send the image file in a request: +## Benchmarking with Apache Bench - curl -X POST -F file=@cat_pic.jpeg http://localhost:5000/predict +Remove Docker: + + docker stop pytorch_flask + docker rm pytorch_flask + +Run Benchmark Code: + docker run --name pytorch_flask -p 5000:5000 -d pytorch_flask:v0 python benchmark.py + +Benchmarking: + + ab -n 100 -c 10 http://localhost:5000/predict ## License The mighty MIT license. Please check `LICENSE` for more details. + diff --git a/app.py b/app.py index 2512017..74254e2 100644 --- a/app.py +++ b/app.py @@ -1,6 +1,6 @@ import io import json - +import os from torchvision import models import torchvision.transforms as transforms from PIL import Image @@ -42,4 +42,5 @@ def predict(): if __name__ == '__main__': - app.run() + port = int(os.environ.get('PORT', 5000)) + app.run(debug=True, host='0.0.0.0', port=port) diff --git a/benchmark.py b/benchmark.py new file mode 100644 index 0000000..c278f53 --- /dev/null +++ b/benchmark.py @@ -0,0 +1,49 @@ +import io +import json +import os +from torchvision import models +import torchvision.transforms as transforms +from PIL import Image +from flask import Flask, jsonify, request + + +app = Flask(__name__) +imagenet_class_index = json.load(open('imagenet_class_index.json')) +model = models.densenet121(pretrained=True) +model.eval() + + +def transform_image(image_bytes): + my_transforms = transforms.Compose([transforms.Resize(255), + transforms.CenterCrop(224), + transforms.ToTensor(), + transforms.Normalize( + [0.485, 0.456, 0.406], + [0.229, 0.224, 0.225])]) + image = Image.open(io.BytesIO(image_bytes)) + return my_transforms(image).unsqueeze(0) + + +def get_prediction(image_bytes): + tensor = transform_image(image_bytes=image_bytes) + outputs = model.forward(tensor) + _, y_hat = outputs.max(1) + predicted_idx = str(y_hat.item()) + return imagenet_class_index[predicted_idx] + + +@app.route('/predict') +def predict(): + #with open('/docker_demo/cat_pic.jpeg', 'rb') as fd: + byteImgIO = io.BytesIO() + byteImg = Image.open("/docker_demo/cat_pic.jpeg") + byteImg.save(byteImgIO, "jpeg") + byteImgIO.seek(0) + img_bytes = byteImgIO.read() + class_id, class_name = get_prediction(image_bytes=img_bytes) + return jsonify({'class_id': class_id, 'class_name': class_name}) + + +if __name__ == '__main__': + port = int(os.environ.get('PORT', 5000)) + app.run(debug=True, host='0.0.0.0', port=port) diff --git a/cat_pic.jpeg b/cat_pic.jpeg new file mode 100644 index 0000000..ccc6afb Binary files /dev/null and b/cat_pic.jpeg differ diff --git a/densenet121-a639ec97.pth b/densenet121-a639ec97.pth new file mode 100644 index 0000000..e2def43 Binary files /dev/null and b/densenet121-a639ec97.pth differ diff --git a/requirements.txt b/requirements.txt index 6ad8b7e..95d354c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ -Flask==1.0.3 -torchvision==0.3.0 +Flask==2.0.3 +torch==1.7.1 +torchvision==0.8.2