Deploying an ML Model inside a Docker Container using Flask.

Piyush Mehta
4 min readMay 29, 2021
Our Final Product, A Flask Web App which uses ML to predict salary based on work experience.

Before we start building our project, I’d highly suggest you quickly go through my previous story on docker containers.

  • In this project, we’ll build an interactive web-app using Flask.
  • This web-app will help us predict salary if a person joins a company ‘X’ based on the company’s past data.

Note: Salary is directly (and only) related to the past work experience of the candidate.

  • We’ll use python programming language to train our machine learning model.
  • Flask will be responsible for our front end as well as backend.
  • Flask app fetches the required data entered by user, processes it in the backend and displays the desired output.
  • The model along with Flask app will be shiped to a docker container later to be hosted.

Training the Model:

  • To train the model, we’ll be using this dataset.
  • Dataset contains two attributes, year’s of experience and the salary of the employee.
  • We’ll be using Linear Regression, to train our model.
  • The overall goal is to find a suitable value of ‘a’ and ‘b’ that fits well in the equation, Y = a + bX.

where, X (Experience) is the independent variable and plotted along the x-axis

Y (Salary) is the dependent variable and plotted along the y-axis

b is the slope of the line, and a is the intercept (the value of y when x = 0).

https://github.com/pewxh/salary-prediction-flask/tree/main/MODEL

Note: at line 16, we dump our model into a binary file named ‘model.x’, this allows us to reuse or even share the model among various programs which can use this pre-trained model directly.

Creating Flask App:

  • Our Flask app provides 2 endpoints.
  • ‘/’ is used to render the main HTML page of our webapp, this page gets the input from user and forwards it to the other route ‘/predict’ using GET call.
  • ‘/predict’ route fetches the user data(name,experience) incoming from ‘/’ route.
  • The experience (x) with the help of model (which was loaded in Flask from the ‘model.x’ file), is used to predict the salary.

The model.predict([[x]]) uses the calculated slope ‘b’ and intercept ‘c’ along with passed experience ‘x’, to predict the salary ‘y’.

Y = a + bX

  • After the successful calculation of the salary, a new HTML page is rendered displaying our predicted salary.
  • In case of any error/exception, Flask handles it and an error page is rendered.

Runing the app inside Docker:

  • Now since the Model is already trained and Flask app is ready, its time to deploy them to the Docker container.
  • There are number of ways to do so, feel free to choose whatever you like.
  • We will be using Dockerfile to create a docker image and git to store and copy the model and Flask code.
  • Build the image and then run the container to deploy the web-app.

After the image is built, we can store it in a remote repository (docker hub), this allows us to easily download the image when ever required.

docker run -d -p 5000:5000 pewxh/salary-prediction-flask

The above commands launches a new detached container with port number 5000 published (to the outside world using port number 5000 of the host) from the image pewxh/salary-prediction-flask:latest.

--

--