Deploy PyTorch NLP – Flask, Docker

Vivek Sasikumar
3 min readMar 10, 2020

--

After training of your PyTorch model, how would one deploy the model in production?

Following is a simple implementation using Python Flask library to connect to a RESTful static HTML page. Afterwards, I convert it into an executable Docker file which can be deployed on AWS ECS (Elastic Container Service). To update the trained model (after training with new data or better optimization), it is as easy as updating the new weights of the model into the container.

Going from the previous blog on translation tool, I am just going to use the trained model from this blog to export model weights.

Export model & information

Now let us create an executable .py file with Flask which accesses the above model weights as well as the dictionary. You have to save the class files in the same folder to import Lang, EncoderRNN, AttnDecoderRNN etc. You can find the class and details in the translation tool blog.

Once this is done, you can go to Powershell or Bash or other command prompt to run the REST API file in local machine.

Add home.html, result.html under folder “templates” and style.css under folder “static”.

In command prompt, go to the folder with translation.py and then enter “python translation.py”

This will initiate the REST API Flask file. Open browser and enter ‘http://0.0.0.0:8000/’ or ‘http://127.0.0.1:8000/’ or ‘http://localhost:8000/’. The following page should show up.

When you enter german (the language we trained), click predict.

Here, the model has just been trained for 50000 epochs (with Google Colab GPU). Ideally, we should be training with a lot more epochs and optimize on learning rate too.

We have successfully deployed the model on your local machine.

Now we move on to dockerizing.

Requirements for python3.6 with needed libraries such PyTorch etc.

Create ‘Dockerfile’ with no extensions and exact format (capital D is a must)

The following is how the folder with Dockerfile should look like. (I am using VSCode on Windows)

Now open command prompt and go to the folder where Dockerfile is saved. Enter “docker build -t translation .

This will pull Python3.6 image from dockerhub, install all dependencies in requirements.txt. This will take some time since PyTorch download and installation takes a while. (note that I did not install Cuda toolkit since we are not going to use GPU).

Also I forgot to add the class files Lang, EncoderRNN, AttnDecoderRNN in the folder. You will have to add it or copy the class code into translation.py.

Now you can deploy docker container on your local machine by typing

docker run -it — name nmt1 -p 8000:8001 translation

This will create a docker container with RESTful Flask file running in it. Open browser and enter ‘http://0.0.0.0:8000/’ or ‘http://127.0.0.1:8000/’ or ‘http://localhost:8000/’.

You can run more containers in parallel.

I will add AWS ECS deployment with Elastic Load Balancer and managing multiple containers in the coming weeks. Also will add deployment of model in Kubernetes and AWS EKS.

--

--