Using Docker for Model Deployment: A Step-by-Step Guide

Introduction

Training a machine learning (ML) model is an important step, but the real challenge lies in deploying it consistently across different environments. Docker is one of the most effective tools for deploying ML models because it packages your application and dependencies into portable containers. This ensures your model runs smoothly, whether on your laptop, in testing, or in production.

In this article, we’ll explore why Docker is important for ML deployment, how to containerize ML models, and the advantages and real-world applications of Dockerized deployments.

Why Use Docker for ML Model Deployment?

Docker provides several benefits for ML engineers and data scientists:

• Portability: Run anywhere without dependency issues.

• Consistency: Same environment across dev, test, and production.

• Scalability: Easy to replicate containers for handling high traffic.

• Integration: Works seamlessly with CI/CD pipelines and cloud providers.

Prerequisites

Before getting started, ensure you have:

• Docker installed (docker --version to verify).

• A trained ML model (Scikit-learn, TensorFlow, or PyTorch).

• Basic knowledge of Flask or FastAPI for serving ML models.

Step 1: Train and Save the Model

Here’s an example with Scikit-learn:

# train_model.py
import joblib
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

iris = load_iris()
X, y = iris.data, iris.target
model = RandomForestClassifier()
model.fit(X, y)

joblib.dump(model, "iris_model.pkl")
print("Model saved as iris_model.pkl")

Step 2: Create a Flask API for the Model

# app.py
from flask import Flask, request, jsonify
import joblib
import numpy as np

app = Flask(__name__)
model = joblib.load("iris_model.pkl")

@app.route("/")
def home():
    return "ML Model with Docker is running!"

@app.route("/predict", methods=["POST"])
def predict():
    data = request.get_json(force=True)
    prediction = model.predict([np.array(data["features"])])
    return jsonify({"prediction": int(prediction[0])})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)

Step 3: Write a Dockerfile

# Use Python base image
FROM python:3.9-slim

# Set working directory
WORKDIR /app

# Copy dependencies
COPY requirements.txt requirements.txt
RUN pip install -r requirements.txt

# Copy project files
COPY . .

# Expose app port
EXPOSE 5000

# Run the app
CMD ["python", "app.py"]

requirements.txt:

flask
scikit-learn
joblib
numpy

Step 4: Build and Run the Docker Container

# Build Docker image
docker build -t ml-flask-app .

# Run container
docker run -p 5000:5000 ml-flask-app

Visit: http://127.0.0.1:5000/

Send a POST request:

curl -X POST http://127.0.0.1:5000/predict \
-H "Content-Type: application/json" \
-d '{"features":[5.1,3.5,1.4,0.2]}'

Output:

{"prediction": 0}

Step 5: Deploy to Cloud

Once containerized, your ML model can be deployed on:

• AWS ECS/EKS

• Google Cloud Run or GKE

• Azure AKS or Container Instances

• Heroku (with Docker support)

Advantages of Docker Deployment

• Consistent environment across machines.

• Easy to scale with Kubernetes.

• Cloud provider agnostic.

• Works well in CI/CD pipelines.

Limitations

• Requires some learning curve.

• Overhead compared to local execution.

• Scaling needs orchestration tools like Kubernetes.

Real-World Applications

• Healthcare: Deploy ML models for medical image analysis.

• Finance: Fraud detection microservices.

• E-commerce: Recommendation engines.

• AI APIs: NLP or CV models as REST endpoints.

Conclusion

Docker is a powerful tool for deploying ML models in a portable, reproducible, and scalable way. By containerizing your model with Flask or FastAPI, you can ensure smooth deployments across different environments and easily integrate into production systems.

For ML engineers looking to go beyond Jupyter notebooks, Docker is an essential step toward production-ready ML applications.