Deep Learning: Loading and Operationalizing Our Model

Image: Shutterstock

What’s the need for building neural networks/modeld if we can’t use them? In this article, we shall look at how to operationalize a neural network.

This is a continuation from my previous article about creating an image classifier using cifar-10. You can find it here.

In the previous article, we trained and saved our model. In this article, we shall load the model and operationalize using python flask.

We have three main steps to complete:

• load our model
• write the prediction function
• wrap to a flask app

Load Our Model

In the last article, we saved our model with the name checkpoint.pth. To load it, we make use of the torch.load function. And then we reconstruct our model.

Please note that for the reconstruction to work, we need our Net class, so we will bring it over.

# let's do our importa here
import torch
from PIL import Image
import torchvision
import torchvision.transforms as transforms
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

# the cifar-10 class labels
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
# our nn
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.conv3 = nn.Conv2d(64, 128, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 2 * 2, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 32)
        self.fc4 = nn.Linear(32, 10)
        self.dropout1 = nn.Dropout(p=0.2, inplace=False)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.dropout1(x)
        x = self.pool(F.relu(self.conv2(x)))
        x = self.dropout1(x)
        x = self.pool(F.relu(self.conv3(x)))
        x = self.dropout1(x)
        x = x.view(-1, 128 * 2 * 2)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x) #output layer
    return x

Let’s check if our machine has GPU, if not, we use CPU

device = 'cuda' if torch.cuda.is_available() else 'cpu'

Now let’s get to the function to load the model

# function to load and reconstruct the model
def load_checkpoint(filepath):
    # we do map_location=device so that in case the model was trained with GPU and then we're trying to load it on a machine that has but CPU, it should still work
    checkpoint = torch.load(filepath, map_location=device)
    model = checkpoint['model']
    model.load_state_dict(checkpoint['state_dict'])
    for parameter in model.parameters():
        parameter.requires_grad = False
    
    # put our model to evaluation mode
    model.eval()
    return model

Prediction Function

Now that we have our code snippets to load and reconstruct our model, let’s write a function to do prediction. The function takes the image path as argument and returns the predicted label.

def do_predcition(filepath):
    pil_image = Image.open(filepath)

    # transform the image to a tensor and resize to fit our trained model size  
    img_loader = transforms.Compose([
    transforms.Resize((32,32)),
    transforms.ToTensor()])
    ts_image = img_loader(pil_image).float()
    ts_image.unsqueeze_(0)
    outputs = model(ts_image.to(device))
    _, predicted = torch.max(outputs.data, 1)
    
    # the prediction is a list, so we have to get the first element    
    return classes[predicted[0]]

At this point, we have everything that has to do with neural networks. What we have left is to wrap the functions to a flask app.

Flask App

Let’s get the imports that are required for our flask app.

# imports specific to the flask app
from flask import Flask, render_template, request
from werkzeug.utils import secure_filename
import os

Basically, we have a simple folder structure as:

root
--templates
----index.html
--app.py

The index.html is a simple form that allows the user to select an image. This part can even be left out if you want to use postman to test the model endpoint. The index.html page is quite simple:

<--index.html-->
<html>
    <body>
        <form action = "/uploader" method = "POST" enctype = "multipart/form-data">
            <input type = "file" name = "file" />
            <input type = "submit"/>
        </form>
    </body>
</html>

Now let’s get to app.py code. We basically have two routes, and then bringing in out prediction code.

# bring all imports here
app = Flask(__name__)
# bring Net class and the other functions and codes here
model = load_checkpoint('checkpoint.pth')
model.to(device)
@app.route('/')
def upload_file_page():
    return render_template('index.html')
@app.route('/uploader', methods = ['GET', 'POST'])
def upload_file():
    image_path = ''
    if request.method == 'POST':
        f = request.files['file']
        f.save(secure_filename(f.filename))
        image_path = "./" + secure_filename(f.filename)
        # call our prediction function
        category = do_predcition(image_path)
        # delete the uploaded file
        os.remove(image_path)
        return category
    # you can hangle the get request here anyhow you'll like
if __name__ == '__main__':
# uncomement this line if you want to run inproduction
#   app.run(host='0.0.0.0', port=80)
    app.run(debug = True)

You can then run your app with python app.py

If you’re running locally, then you can find the app at http://127.0.0.1:5000/

If you run on your cloud server, you can access via your server public ip(like in my case), of course after setting the right config on the app.py code above.

And that’s it.

Thanks for your time! Hope you enjoyed it!

References

• Torch load
• Flask documentation