Loading a machine learning model in docker

Arjun Chauhan
4 min readMay 27, 2021

In today’s blog I would be doing a simple thing. I will create a docker container and run a machine learning model in it.

So machine learning models are quite memory consuming although the model and the dataset I have used are pretty small but real world applications can be much more memory extensive.

So I will be training the model in my local system and then move the saved model in docker container.

This has multiple benefits. We can use the container to create a new docker image which can be used to create multiple containers who are capable of predicting.

The steps

  1. I would be using the RHEL 8 system as by BaseOS. So firstly we need to download docker. For this we use

yum install docker

Docker is just like any other service which needs to be started or enabled. To start docker and check whether it is running or not we use

systemctl start docker

systemctl status docker

2. Now since the services are up and running we can now create a docker container. For this I am using the centos:latest image.

The command to run the container is

docker run -dit <container_name> <image_name>:<image_tag> /bin/bash

As you can see the container is up and running. To view the status of the container we can use

docker ps

  • dit used in the above command tells the docker that we need a detached interactive terminal for our container.

3. Now we have the bare minimum container requirement done. So now moving on to the python script to train the model. The dataset used here is a small artificial dataset which uses the experience to predict the salary of a person.

It is very easy to see that this is a simple linear regression problem. In the below file you see I have read this dataset and used python’s ML library scikit learn to fit the dataset and save the model in a pickle file.

4. We would now create a python script that is going to predict the salary of a person using his years of experience as input

5. Now all that is left is to copy this file in the docker container and set up the python environment. To copy file to the docker container we use

docker cp <source_location> <container_name>:<destination>

After transferring the file we can attach ourselves to the container

6. Now we need to install the python3 and scikit-learn libraries in our docker container

7. Now finally we can run the python script and get the output as the salary when we get the number of years as the input from the user

This was a very simple task however we can build around this. For example we can create a kubernetes cluster of this container and then use a webserver to host it. However this is something to dwell upon some other day

Thanks

--

--