" How to Deploy a Machine Learning Model Using Flask on Ubuntu 24.04 LTS "

Deploying a machine learning model is a crucial step to making it accessible to end-users. Flask, a lightweight web framework for Python, is popular for deploying machine learning models due to its simplicity and flexibility. When deploying on an Ubuntu 24.04, you can leverage the power of GPUs to accelerate inference, especially for deep learning models.

Flask is a micro web framework written in Python.
In this Tutorial you will Learn  How to Deploy a Machine Learning Model Using Flask on Ubuntu 24.04
_____________________________________________________________________________________________________________
Testing Environment: Ubuntu 24.04.1 LTS       Hostname -          ip Address -
_____________________________________________________________________________________________________________
apt update ; apt install build-essential net-tools curl git software-properties-common

apt install python3 python3-pip python3-virtualenv python3.12-venv -y
python3 -m venv ml-env
source ml-env/bin/activate
pip install --upgrade pip
pip install flask gunicorn torch torchvision tensorflow numpy pandas
nano train_model.py
import torch
import torch.nn as nn
import torch.optim as optim

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(4, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 1)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        # Using sigmoid for a binary classification output
        x = torch.sigmoid(self.fc3(x))
        return x

# Create the model
model = SimpleModel()

# Define a binary cross-entropy loss and Adam optimizer
criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Dummy training data:
data = torch.rand(100, 4)
labels = torch.randint(0, 2, (100, 1), dtype=torch.float32)

# Train for 10 epochs
for epoch in range(10):
    optimizer.zero_grad()
    outputs = model(data)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    if (epoch+1) % 2 == 0:
        print(f"Epoch [{epoch+1}/10], Loss: {loss.item():.4f}")

# Switch model to evaluation mode (optional but good practice)
model.eval()

# Save the model's state dict only (recommended approach)
torch.save(model.state_dict(), "model_weights.pth")
print("Model weights saved to 'model_weights.pth'")

python3 train_model.py
nano app.py
import torch
import torch.nn as nn
import numpy as np
from flask import Flask, request, jsonify

app = Flask(__name__)

###############################################
# 1. Define the EXACT SAME PyTorch architecture
###############################################
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        # Must match your training script's layers:
        self.fc1 = nn.Linear(4, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 1)  # Single output (e.g. for binary classification)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.sigmoid(self.fc3(x))
        return x

###########################################################
# 2. Load the PyTorch model state dict (must match architecture)
###########################################################
pytorch_model = SimpleModel()
pytorch_model.load_state_dict(torch.load("model_weights.pth"))
pytorch_model.eval()

###########################################################
# 3. Flask endpoint for PyTorch inference
###########################################################
@app.route("/predict", methods=["POST"])
def predict():
    data = request.get_json()

    # Convert input to a numpy array of shape [1, 4] since the model expects 4 features
    input_data = np.array(data["input"]).reshape(1, 4)
    
    # PyTorch inference
    tensor_input = torch.tensor(input_data, dtype=torch.float32)
    pytorch_output = pytorch_model(tensor_input).detach().numpy().tolist()

    return jsonify({
        "pytorch": pytorch_output
    })

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

python3 app.py
curl -X POST -H "Content-Type: application/json" -d '{"input":[0.1, 0.2, 0.3, 0.4]}' http://localhost:5000/predict