Build a Machine Learning Model API with FastAPI
Imagine you’ve spent weeks training a machine learning model, tweaking hyperparameters, and watching accuracy curves climb. But now it sits in a .pkl file on your laptop, useless to anyone else. That’s the gap between a cool experiment and a real product: deployment. The fastest, most modern way to bridge that gap is FastAPI, a framework that lets you wrap your model in a production-ready API with minimal boilerplate and maximum speed.
In this guide, you’ll build a complete Machine Learning Model API using FastAPI and Scikit-Learn. By the end, you’ll have a working endpoint that accepts input data, runs a prediction, and returns JSON—ready to be consumed by a frontend, mobile app, or another service. Let’s get your model out of the notebook and into the world.
Why FastAPI for Machine Learning?
FastAPI isn’t just another Python web framework. It’s built on Starlette for the HTTP layer and Pydantic for data validation, giving you three critical advantages for ML:
- Performance: FastAPI is among the fastest Python frameworks available, thanks to asynchronous support and native JSON handling.
- Automatic Validation: Pydantic models automatically validate and serialize incoming data, so you never have to write manual type checks.
- Interactive Docs: FastAPI generates Swagger UI and ReDoc documentation automatically, letting you test endpoints instantly without writing extra test code.
For ML models that might need to serve thousands of requests per second, these features are non-negotiable.
Step 1: Set Up Your Environment
First, create a virtual environment and install the required packages:
python -m venv ml-api-env
source ml-api-env/bin/activate # On Windows: ml-api-env\Scripts\activate
pip install fastapi uvicorn scikit-learn pydantic
-
fastapi: The core framework. -
uvicorn: The ASGI server that runs FastAPI. -
scikit-learn: For our sample ML model. -
pydantic: For data validation (already included with FastAPI).
Step 2: Create a Dummy ML Model
You don’t need a real trained model to start. Let’s create a simple house price predictor using Scikit-Learn. This will be saved as a .pkl file so FastAPI can load it later.
# train_model.py
from sklearn.linear_model import LinearRegression
import pickle
# Sample data: square feet -> price
X = [[500], [600], [700], [800], [900], [1000]]
y = [100000, 120000, 140000, 160000, 180000, 200000]
model = LinearRegression()
model.fit(X, y)
# Save the model
with open("house_price_model.pkl", "wb") as f:
pickle.dump(model, f)
print("Model trained and saved as 'house_price_model.pkl'")
Run this script once to generate the model file:
python train_model.py
Step 3: Build the FastAPI Application
Now, create your main API file: api.py. This script will:
- Load the saved model.
- Define input validation with Pydantic.
- Create a
/predictendpoint that returns predictions.
# api.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import pickle
# Initialize FastAPI app
app = FastAPI(title="House Price Prediction API", version="1.0.0")
# Load the trained model
model = None
@app.on_event("startup")
def load_model():
with open("house_price_model.pkl", "rb") as f:
model = pickle.load(f)
# Define input data model
class HouseInput(BaseModel):
square_feet: int
# Prediction endpoint
@app.post("/predict")
def predict_price(input_data: HouseInput):
if model is None:
raise HTTPException(status_code=500, detail="Model not loaded")
# Prepare data for model
X = [[input_data.square_feet]]
prediction = model.predict(X)[0]
return {
"square_feet": input_data.square_feet,
"predicted_price": round(prediction, 2)
}
# Health check endpoint
@app.get("/")
def read_root():
return {"message": "Welcome to the House Price Prediction API"}
Step 4: Run and Test Your API
Start the server with Uvicorn:
uvicorn api:app --reload
The --reload flag enables auto-reloading on code changes, perfect for development.
Once running, open your browser to http://localhost:8000. You’ll see:
- The root endpoint (
/) with a welcome message. -
Swagger UI at
http://localhost:8000/docs, where you can test the/predictendpoint interactively.
Testing with a Real Request
Use curl to send a prediction request:
curl -X POST "http://localhost:8000/predict" \
-H "Content-Type: application/json" \
-d '{"square_feet": 750}'
Expected response:
{
"square_feet": 750,
"predicted_price": 150000.0
}
You can also test it from Python using requests:
import requests
response = requests.post(
"http://localhost:8000/predict",
json={"square_feet": 750}
)
print(response.json())
Step 5: Production Tips for Real ML APIs
While this example uses a simple model and local file loading, real-world deployments require more robustness. Here’s what to add next:
| Feature | Why It Matters |
|---|---|
| Async Functions | Prevent blocking when the model takes time to compute [6] |
| Logging & Monitoring | Track errors, latency, and prediction volume [7] |
| Dockerization | Ensure consistent environments across dev and production [8] |
| Model Storage in AWS S3 | Avoid storing large models in your app directory [7] |
| Background Tasks | Run heavy preprocessing/postprocessing without blocking the response [7] |
| Unit Tests | Validate that your API returns correct predictions for known inputs [7] |
For production, consider using multistage Docker builds to reduce image size and storing models in cloud storage like AWS S3 instead of local files [7][8].
Your Model is Now Live
You’ve just built a fully functional ML API that can be consumed by any client. No more Jupyter notebooks—your model is now a service.
What’s Next?
- Replace the dummy model with your own trained model (e.g., from TensorFlow, PyTorch, or XGBoost).
- Add input preprocessing (e.g., scaling, encoding) inside the endpoint [9].
- Implement authentication and rate limiting for public APIs.
- Deploy to cloud platforms like AWS, Google Cloud, or Azure using Docker [8].
Call to Action
Try it today: Copy the code above, run it on your machine, and send a prediction request. Then, share your API endpoint in the comments or on Twitter—tag me so I can see what you built!
If you found this helpful, check out my other posts on deploying ML models with Docker and building async ML pipelines. Let’s keep pushing ML from research to real-world impact. 🚀
If you found this helpful, consider buying me a coffee ☕ — it keeps these articles coming!
Also check out my AI tools collection: AI 次元世界 — free AI tools for developers.
💡 Related: **Content Creator Ultimate Bundle (Save 33%)* — $29.99*
Top comments (0)