Deploying a Neural Network with FastAPI, Docker, and Kubernetes

In this guide, we will deploy a neural network using FastAPI, containerize it with Docker, and orchestrate the deployment with Kubernetes.

1. Train and Save the Model

We first create a simple neural network and save it as model.h5.

import tensorflow as tf
from tensorflow import keras
import numpy as np

# Create and train the model
model = keras.Sequential([
    keras.layers.Dense(128, activation='relu', input_shape=(10,)),
    keras.layers.Dense(64, activation='relu'),
    keras.layers.Dense(1, activation='sigmoid')
])

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# Save the model
model.save("model.h5")
print("Model saved as 'model.h5'")

2. Create FastAPI Application

Next, we build a FastAPI app to serve predictions.

from fastapi import FastAPI
from pydantic import BaseModel
import tensorflow as tf
import numpy as np

app = FastAPI()
model = tf.keras.models.load_model("model.h5")

class InputData(BaseModel):
    input: list

@app.post("/predict")
async def predict(data: InputData):
    input_data = np.array([data.input])
    prediction = model.predict(input_data).tolist()
    return {"prediction": prediction}

3. Dockerizing the Application

FROM python:3.9
WORKDIR /app
COPY model.h5 /app/
COPY app.py /app/
RUN pip install fastapi uvicorn tensorflow numpy
EXPOSE 8000
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

4. Kubernetes Deployment

Now, we deploy the application on Kubernetes using deployment.yaml and service.yaml.

Deployment YAML

apiVersion: apps/v1
kind: Deployment
metadata:
  name: fastapi-neural-net
spec:
  replicas: 2
  selector:
    matchLabels:
      app: neural-net
  template:
    metadata:
      labels:
        app: neural-net
    spec:
      containers:
      - name: neural-net
        image: your-dockerhub-username/fastapi-neural-net:latest
        ports:
        - containerPort: 8000

Service YAML

apiVersion: v1
kind: Service
metadata:
  name: neural-net-service
spec:
  selector:
    app: neural-net
  ports:
    - protocol: TCP
      port: 80
      targetPort: 8000
  type: LoadBalancer

5. Deploy and Access

Run the following commands:

kubectl apply -f deployment.yaml
kubectl apply -f service.yaml
kubectl get services

Conclusion

We have successfully deployed a neural network using FastAPI, Docker, and Kubernetes. This setup ensures scalability and efficiency in machine learning model deployment.