Deploy Scikit-Learn & TensorFlow Models with FastAPI

In the previous lesson, we explored how to deploy machine learning models using FastAPI. We learned how to set up a basic FastAPI app, integrate a simple model, and handle API requests. Now, we’ll take it a step further by deploying more advanced models trained with Scikit-Learn and TensorFlow. These tools are widely used in data science, and combining them with FastAPI allows us to build powerful, real-time prediction systems.

Deploying a Scikit-Learn Model with FastAPI

I recently worked on a project where I needed to deploy a Scikit-Learn model for predicting house prices. The model was trained on a dataset that included features like square footage, number of bedrooms, and location. My goal was to create an API that could take these features as input and return the predicted price in real-time. FastAPI was the perfect choice because of its speed and ease of use.

To achieve this, I first saved the trained Scikit-Learn model using joblib. Then, I loaded the model into a FastAPI app and created an endpoint that accepted input data, processed it, and returned the prediction. The entire process was straightforward, and the API performed well under heavy load. This experience showed me how powerful FastAPI can be for serving machine learning models.

Steps to Serve Scikit-Learn Models with FastAPI

1. Save the Trained Model: Use joblib or pickle to save your Scikit-Learn model after training. For example:

import joblib
joblib.dump(model, 'house_price_model.pkl')

2. Load the Model in FastAPI: Create a FastAPI app and load the saved model. Here’s how:

from fastapi import FastAPI
import joblib

app = FastAPI()
model = joblib.load('house_price_model.pkl')

3. Create an Endpoint: Define an endpoint that accepts input data and returns predictions. For example:

from pydantic import BaseModel

class HouseFeatures(BaseModel):
    sqft: float
    bedrooms: int
    location: str

@app.post("/predict")
def predict(features: HouseFeatures):
    input_data = [[features.sqft, features.bedrooms, features.location]]
    prediction = model.predict(input_data)
    return {"predicted_price": prediction[0]}

4. Test the API: Run the FastAPI app and test the endpoint using tools like Postman or directly in the browser.

5. Optimize for Performance: Use FastAPI’s async capabilities to handle multiple requests efficiently. For example:

@app.post("/predict")
async def predict(features: HouseFeatures):
    input_data = [[features.sqft, features.bedrooms, features.location]]
    prediction = model.predict(input_data)
    return {"predicted_price": prediction[0]}

TensorFlow Model Inference with FastAPI

In another project, I needed to deploy a TensorFlow model for image classification. The model was trained to classify images of cats and dogs. I used FastAPI to create an API that could accept an image, process it, and return the predicted class. FastAPI’s support for file uploads made this task easy.

To implement this, I saved the TensorFlow model using the SavedModel format. Then, I loaded the model into the FastAPI app and created an endpoint that accepted image uploads. The API processed the image, ran it through the model, and returned the result. This approach worked well, and the API was able to handle multiple requests without any issues.

Steps to Serve TensorFlow Models with FastAPI

1. Save the TensorFlow Model: Use the SavedModel format to save your trained model. For example:

model.save('image_classifier_model')

2. Load the Model in FastAPI: Load the saved model in your FastAPI app:

import tensorflow as tf
from fastapi import FastAPI, File, UploadFile

app = FastAPI()
model = tf.keras.models.load_model('image_classifier_model')

3. Create an Endpoint: Define an endpoint that accepts image uploads and returns predictions:

@app.post("/classify")
async def classify(image: UploadFile = File(...)):
    image_data = await image.read()
    processed_image = preprocess_image(image_data)
    prediction = model.predict(processed_image)
    return {"predicted_class": "cat" if prediction[0] > 0.5 else "dog"}

4. Test the API: Run the app and test the endpoint by uploading images.

5. Optimize for Performance: Use FastAPI’s async features to handle multiple image uploads efficiently.

Optimizing API Responses for Real-Time Predictions

When deploying machine learning models, response time is critical. Users expect quick results, especially in real-time applications. To optimize performance, I used FastAPI’s async capabilities and ensured that the models were loaded only once during startup. This reduced latency and improved the overall user experience.

For example, in the house price prediction app, I used async functions to handle multiple requests simultaneously. This allowed the API to serve more users without slowing down. Similarly, in the image classification app, I optimized the image preprocessing step to reduce the time taken for each prediction.

Conclusion

In this lesson, we learned how to deploy Scikit-Learn and TensorFlow models using FastAPI. We covered the steps to save and load models, create API endpoints, and optimize performance for real-time predictions. FastAPI’s speed and simplicity make it an excellent choice for serving machine learning models.

If you’re ready to take your skills to the next level, don’t miss the next lesson where we’ll connect FastAPI with PostgreSQL for data science applications. This will allow you to store and retrieve data efficiently, making your apps even more powerful.