Forem: James Casia

Deploying Pytorch Models with Flask & Amazon EC2

James Casia — Mon, 23 Jan 2023 10:44:51 +0000

So you have finally trained and built your machine learning model. It took you days on end doing research, training and fine-tuning to achieve perfection. But in the end, a model is only as good as the good it causes for the people who use it. Thus deploying a machine learning model is an essential step as it allows one to put into action and see its real-world impact on the people it was designed to help.

In this article, we create a very basic API for our pytorch model. We use Flask to build the REST API and AWS EC2 as our deployment platform of choice. This article breaks down the task into four major steps.

Setup & Installing Dependencies

Creating & Using Virtual Environment(optional)

We will be using virtualenv for an isolated python environment. Run the following commands to install, create, and activate our virtual environment.

$ python -m pip install --user virtualenv
$ python -m venv env
$ source env/bin/activate

Installing Requirements

Create a requirements.txt file with the following content.

//requirements.txt

torch==1.13.1
torchvision==0.14.1
Flask==2.2.2 
pandas==1.5.2 
tqdm==4.64.1
matplotlib==3.6.3
gunicorn==20.1.0

Then simply run the command pip install -r requirements.txt to install the dependencies detailed in the requirements.txt file

Writing the code

Project structure

We follow the standard Flask project file structure. We have a static folder that contains our static assets such as CSS and images. We also have a templates directory that contains our HTML templates. Our main directory is app and it contains two files, __init__.py and main.py . For our purpose, we leave __init__.py empty and write the routes and other relevant code in main.py.

app/
    __init__.py
    main.py
static/
    /images
        /resnet/
templates/
requirements.txt
run.sh

Importing dependencies

We import the essential dependencies.

# main.py
from torchvision.models import resnet50, ResNet50_Weights
from flask import Flask, request ,render_template, url_for
from werkzeug.utils import secure_filename 
import urllib.request,io 
from PIL import Image 
import os

Writing the create_app function

In main.py, we write a create_app function that returns a Flask object.

# main.py
def create_app():
        app = Flask('image-classifier')

        ...
        ...

        return app

Initializing the Pytorch Model

For this tutorial, we use a pre-trained resnet50 model available from torchvision for convenience. Feel free to replace this with your own model. We then call .eval() to toggle the model to inference mode.

# main.py
def create_app():

    app = Flask('image-classifier') 

    weights = ResNet50_Weights.DEFAULT
    model = resnet50(weights=weights)
    model.eval()
    preprocess = weights.transforms()

For this model, we create a preprocess object that is basically a pre-processing transformation step that comes with the resnet50 model. We will use this preprocess object to modify the images we feed into the model.

Defining the routes

In this app, we have three routes, /gui and /classify . The /gui route is where we present the graphical user interface to interact with the API. The /classify route is for the classification API itself for which we can interact with it through a command line interface or a REST inspector application. The classify_gui function contains the code for the GUI interaction with the classification API. It is slightly different from the classify function as it contains an image upload feature.

def create_app():

        ...
        ...

    @app.route('/classify')
    def classify(): 
                ...
                ...

    @app.route('/')   
    def gui():
            ...
                ...

    @app.route('/', methods=['POST'])  
    def classify_gui():
            ...
                ...

    return app

Defining the classification endpoint

Flask allows you to define the routes through the @app.route decorator. In the classify function, we implement what happens when a user performs a GET request in the /classify route. We nest this function in the create_app function for convenience. This allows us to have access to the model without passing it as parameters to classify().

# main.py
def create_app():

        ...
        ...

    @app.route('/classify')
    def classify(): 
        url = io.BytesIO(urllib.request.urlopen(request.get_json()['url']).read())
        img = Image.open(url)

        batch = preprocess(img).unsqueeze(0) 
        prediction = model(batch).squeeze(0).softmax(0)
        class_id = prediction.argmax().item()
        score = prediction[class_id].item()
        category_name = weights.meta["categories"][class_id] 
        return  { 'class': category_name, 'confidence': float("{:.2f}".format(score)) }

In this particular function, we accept a json containing a url value. We parse the url from the request body and load it as a Pillow image. We then do some model-specific pre- and post- processing methods. The resnet50 model by default accepts a batch of images. Since we’re only passing it one image, we perform unsqueeze to make it a batch of one image. After passing the batch to model, we extract the prediction by applying the softmax function and getting the index of the maximum prediction value.

Defining the GUI endpoint

The / route will contain the graphical interface for users to interact with the model. Here is how it will look like.

We use render_template function to return html.

    @app.route('/')   
    def gui():
        return render_template('gui.html')

We create the base.html file. The gui.html will inherit from this file

<!--base.html-->
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta http-equiv="X-UA-Compatible" content="IE=edge">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">  
    <title>ML API</title>
</head>
<body>
{% block container %}{% endblock %}
</body>
</html>

We create the gui.html file.

<!--gui.html-->
{% extends 'base.html' %}

{% block container %}
<div class="home"> 
<h1>Image Classification</h1>
<h3>with Pytorch ResNet50</h3>  
<h2>CLI</h2>
<table border="1">
    <tr>
        <td><strong>curl</strong></td>
        <td>curl --request GET \
            --url BASE_URL/classify \
            --header 'Content-Type: application/json' \
            --data '{
              "url": IMG_URL
          }'</td>
    </tr>
    <tr>
        <td><strong>wget</strong></td>
        <td>wget --quiet \
        --method POST \
        --header 'Content-Type: application/json' \
        --body-data '{\n    "url": IMG_URL \
        --output-document \
        - BASE_URL/classify</td>
    </tr>
     </table>
</div>

<h2>GUI</h2>
<form method="post" enctype="multipart/form-data" >
    {{img_path}}
    <input type="file" name="file" id="img_upload"><br><br>

    <img id="image" height="400px" src={{img_url}}>

    <br>
    <input type="submit" value="Classify">
</form> 

<br>
<div>{{result}}</div>
<script>
    document.getElementById('img_upload').onchange = function ( ) { 
   var src = URL.createObjectURL(this.files[0])
  document.getElementById('image').src = src
}
</script>
{% endblock %}

Then we craft the classify_gui function. This function gets called when the user presses the classify button.

    @app.route('/', methods=['POST'])  
    def classify_gui():
        clear_dir('./static/images/') 
        if 'file' not in request.files:
            return render_template('gui.html', msg='No file found')

        file = request.files['file']
        if file.filename == '':
            return render_template('gui.html', msg='No file selected')

        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            path = f'static/images/{filename}'
            file.save(path)

            img = Image.open(path)    
            batch = preprocess(img).unsqueeze(0) 

            prediction = model(batch).squeeze(0).softmax(0)
            class_id = prediction.argmax().item()
            score = prediction[class_id].item()
            category_name = weights.meta["categories"][class_id]  

            return render_template(
                'gui.html', result = f"{category_name}: {100 * score:.1f}%", 
                img_url = f"{url_for('static', filename=f'/images/{filename}')}"
            )

Code Summary

Here is the whole code for main.py

# main.py
from torchvision.models import resnet50, ResNet50_Weights
from flask import Flask, request ,render_template, url_for
from werkzeug.utils import secure_filename 
import urllib.request,io 
from PIL import Image 
import os 

def create_app():

    app = Flask('image-classifier') 

    weights = ResNet50_Weights.DEFAULT
    model = resnet50(weights=weights)
    model.eval()
    preprocess = weights.transforms()

    @app.route('/classify')
    def classify(): 
        url = io.BytesIO(urllib.request.urlopen(request.get_json()['url']).read())
        img = Image.open(url)

        batch = preprocess(img).unsqueeze(0) 
        prediction = model(batch).squeeze(0).softmax(0)
        class_id = prediction.argmax().item()
        score = prediction[class_id].item()
        category_name = weights.meta["categories"][class_id] 
        return  { 'class': category_name, 'confidence': float("{:.2f}".format(score)) } 

    @app.route('/')   
    def gui():
        return render_template('gui.html')

    @app.route('/', methods=['POST'])  
    def classify_gui():
        clear_dir('./static/images/') 
        if 'file' not in request.files:
            return render_template('gui.html', msg='No file found')

        file = request.files['file']
        if file.filename == '':
            return render_template('gui.html', msg='No file selected')

        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            path = f'static/images/{filename}'
            file.save(path)

            img = Image.open(path)    
            batch = preprocess(img).unsqueeze(0) 

            prediction = model(batch).squeeze(0).softmax(0)
            class_id = prediction.argmax().item()
            score = prediction[class_id].item()
            category_name = weights.meta["categories"][class_id]  

            return render_template(
                'gui.html', result = f"{category_name}: {100 * score:.1f}%", 
                img_url = f"{url_for('static', filename=f'/images/{filename}')}"
            ) 

    return app

def clear_dir(dir_path):
    for file in os.listdir(dir_path):
        os.remove(f'{dir_path}/{file}') 

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in {'png', 'jpg', 'jpeg'}

This code is also available on github.

Localhost testing

At this point, one can already deploy and use the app in our localhost. Just create a run.sh file with the following commands. Setting the hostname is necessary as it allows any other computer besides the host to access our flask app.

# run.sh
source env/bin/activate
export FLASK_APP="app.main:create_app"
export FLASK_RUN_PORT=8080
export FLASK_APP_HOSTNAME="0.0.0.0" 
flask run --host=$FLASK_APP_HOSTNAME

Then run the shell file

$ sh run.sh

To create a request, simply paste this sample request to another terminal window.

$ curl --request GET \
  --url localhost:8080/classify \
  --header 'Content-Type: application/json' \
  --data '{
    "url": "https://external-content.duckduckgo.com/iu/?u=https%3A%2F%2F2.bp.blogspot.com%2F-HB23AuZMkCc%2FUmFMY7DNM8I%2FAAAAAAAAA8Q%2FFD8D50slVJ4%2Fs1600%2FAlbatross-Bird-Pic.jpg&f=1&nofb=1&ipt=0ac75b04d5d023fa362e560745360a7aad5e15c90e2673d887ff40851e71d9f3&ipo=images"
}'

Or launch your favorite API inspector, set the content type header to application/json and supply an object with the url of your chosen image. Voila! You have now created a REST API for your machine learning model!!

In this image, we pass a url of an image and receive a prediction. The app being used here is Insomnia

Deploying to the Cloud

Deploying your API to your localhost is one thing. Deploying your API to the totality of the interwebs is a whole different matter! Suddenly you have to worry about a whole slew of complexities such as DNS, Inbound & Outbound rules, virtual machines and so much more! Good thing as there are a bajillion articles and tutorials on the internet just about that!

Creating EC2 Instance

Since there are a lot of EC2 tutorials we can simply follow this one on youtube, https://www.youtube.com/watch?v=oqHfiRzxunY. Make sure to select Ubuntu Server AMI, and t2.micro as they are available in the free-tier.

Modifying the security group and Inbound rules

We add an inbound rule to allow other computers to communicate with our instance. We select Custom TCP as the type and 8080 as the port range. We set the source to 0.0.0.0/0 to allow any device from any ip address.

Opening the port in the firewall

Run the following commands in the EC2 terminal to open port 8080.

$ sudo ufw allow 8080
$ sudo ufw enable

Getting your code into the EC2 instance

There are two methods that to get your code into the EC2 instance. One way way is through SCP which is a linux command that allows for file transfer to authenticated users. Another quick but dirty way is to create a public git repository, this way you don’t have to deal with linux complexities. In this article, we’ll simply use the repository route. We simply create a git repository containing the code, publish it, and clone the public repository on the EC2 instance.

Deploying

Run the following commands to download pip and install the necessary libraries.

$ sudo apt update
$ sudo apt install python3-pip
$ pip install -r requirements.txt

Then simply run run.sh again and type in http://PUBLIC_IP_ADDRESS:8080 to any web browser. Use your EC2 instance’s public ip address. Remember to use http instead of https. Voila! You can now access your machine learning API in any part of the world.

Ensuring scalability and security (Part 2 Coming Soon)

We have deployed our Flask app but the Flask team themselves do not recommend deploying bare Flask apps to production. This is because Flask does not come with built-in support for load balancing, security, database management etc. That will be discussed in another article though as this article has gone too long.

Creating Custom Models with Pytorch

James Casia — Wed, 30 Nov 2022 03:14:37 +0000

Pytorch was built with custom models on mind. With just a few lines of code, one can spin up and train a deep learning model in a couple minutes. This is a quick guide to creating typical deep learning models in Pytorch.

Simple Custom Models

To build a custom model, just inherit nn.Module and define the forward function.

class FCN(nn.Module):
    def __init__(self, input_dims, output_dims):
        super(FCN, self).__init__()

        self.model = nn.Sequential(
            nn.Linear(input_dims, 5), 
            nn.LeakyReLU(), 
            nn.Linear(5, output_dims), 
                        nn.Sigmoid()
        )

    def forward(self, X):
        return self.model(X)

And when you instantiate it,

fcn = FCN(10, 1)
fcn

FCN(
(model): Sequential(
(0): Linear(in_features=10, out_features=5, bias=True)
(1): LeakyReLU(negative_slope=0.01)
(2): Linear(in_features=5, out_features=1, bias=True)
(3): Sigmoid()
)
)

Simple as that, you now have a 2-layer neural network! Take note of this pattern as this repeats to more complex models.

A Slightly More Complicated Model

Sometimes, models have repetitive layers, these are called blocks. These are usually found in Convolutional Neural Networks, wherein a block usually contains Convolutional layers, followed by Maxpooling layers. We create these blocks through writing our own custom functions! This time, we'll create a Linear block with BatchNormalization and ReLU. (Don’t worry, we’ll get to CNNs later in this guide)

Simply create a function and return a Sequential model that wraps the layers.

def net_block(input_dim, output_dim):
    return nn.Sequential(
        nn.Linear(input_dim, output_dim),
        nn.BatchNorm1d(output_dim),
        nn.ReLU()
    )

Let’s add a few of these blocks to our model!

class Network(nn.Module):
    def __init__(self, input_dim, hidden_dim_1, hidden_dim_2, output_dim):
        super(Network, self).__init__()

        self.model = nn.Sequential(
            net_block(input_dim, hidden_dim_1),
            net_block(hidden_dim_1, hidden_dim_2),
            net_block(hidden_dim_2, output_dim)
        )

    def forward(self, X):
        return self.model(X)

net = Network(10, 4, 5, 1)
net

Network(
(model): Sequential(
(0): Sequential(
(0): Linear(in_features=10, out_features=4, bias=True)
(1): BatchNorm1d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(1): Sequential(
(0): Linear(in_features=4, out_features=5, bias=True)
(1): BatchNorm1d(5, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(2): Sequential(
(0): Linear(in_features=5, out_features=1, bias=True)
(1): BatchNorm1d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
)
)

We can actually see that we have these repeating blocks. Phew! This saved us from defining these 9 individual layers one by one.

VGG Network

VGG Networks were one of the first deep CNNs ever built! Let’s try to recreate it. First, we define the blocks. VGG has two types of blocks, defined in this paper.

Figure 1. The two types of VGG Blocks: the two-layer(blue, orange) and the three-layer ones(purple, green, red). Image taken from this datasciencecentral article.

Let’s the define the two blocks.

def vgg_block(in_channels, out_channels):  
    return  nn.Sequential(  
        nn.Conv2d(in_channels, out_channels, kernel_size=(3,3), stride=1, padding = (1,1)),
        nn.ReLU(inplace =True), 
        nn.MaxPool2d(kernel_size=(2,2),stride=2)
    )

def vgg_block2(in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=(3,3), stride=1, padding = (1,1)),
        nn.ReLU(inplace =True), 
        nn.Conv2d(out_channels, out_channels, kernel_size=(3,3), stride=1, padding = (1,1)),
        nn.ReLU(inplace =True),  
        nn.MaxPool2d(kernel_size=(2,2),stride=2)
    )

Let’s also define a flatten layer. A flatten layer basically reshapes a tensor into an n-dimensional vector.

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

Let’s define the classification block.

def classifier_block(input_dim, hidden_dim, num_classes):
    return nn.Sequential( 
        Flatten(),
        nn.Linear(input_dim, hidden_dim),
        nn.ReLU(),
        nn.Dropout(0.5),
        nn.Linear(hidden_dim, hidden_dim),
        nn.ReLU(),
        nn.Dropout(0.5),
        nn.Linear(hidden_dim , num_classes),
        nn.Softmax()
    )

Now let’s define the model.

class VGG(nn.Module):
    def __init__(self):
        super(VGG, self).__init__()
        self.model = nn.Sequential( 
            vgg_block(3, 64 ),
            vgg_block(64, 128),
            vgg_block2(128, 256),
            vgg_block2(256, 512), 
            vgg_block2(512, 512),   
            nn.AdaptiveAvgPool2d(output_size=(7, 7)),
            classifier_block(512*7*7, 4096, 1000)
        )

    def forward(self, X):
        return self.model(X)

vgg = VGG()
vgg

VGG(
(model): Sequential(
(0): Sequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): MaxPool2d(kernel_size=(2, 2), stride=2, padding=0, dilation=1, ceil_mode=False)
)
(1): Sequential(
(0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): MaxPool2d(kernel_size=(2, 2), stride=2, padding=0, dilation=1, ceil_mode=False)
)
(2): Sequential(
(0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): ReLU(inplace=True)
(4): MaxPool2d(kernel_size=(2, 2), stride=2, padding=0, dilation=1, ceil_mode=False)
)
(3): Sequential(
(0): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): ReLU(inplace=True)
(4): MaxPool2d(kernel_size=(2, 2), stride=2, padding=0, dilation=1, ceil_mode=False)
)
(4): Sequential(
(0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): ReLU(inplace=True)
(4): MaxPool2d(kernel_size=(2, 2), stride=2, padding=0, dilation=1, ceil_mode=False)
)
(5): AdaptiveAvgPool2d(output_size=(7, 7))
(6): Sequential(
(0): Flatten()
(1): Linear(in_features=25088, out_features=4096, bias=True)
(2): ReLU()
(3): Dropout(p=0.5, inplace=False)
(4): Linear(in_features=4096, out_features=4096, bias=True)
(5): ReLU()
(6): Dropout(p=0.5, inplace=False)
(7): Linear(in_features=4096, out_features=1000, bias=True)
(8): Softmax(dim=None)
)
)
)

We now have recreated the VGG model! How cool is that.

A Complex Model

What we’ve encountered so far are relatively simple networks, they take in the output of the previous layer and transform it. An improvement from this are residual networks, these networks pass the output from an earlier layer onto a layer a few steps ahead forming a “skip connection”. This has been shown to help deep neural networks learn better as gradients can more readily flow through these connections. A few examples of residual networks are ResNets, and UNet.

Residual Networks(ResNets)

Figure 2. ResNet. Image taken from deeplearning.ai Convolutional neural networks course

Residual networks are designed a bit differently to cater to this feature. In these types of networks, we need to have a reference to these layers that have skip connections(which you will see later on).

This time using functions for creating blocks won't cut it anymore. This is because our forward function needs to be customized to accommodate skip connections!

The Residual Block

class ResidualBlock(nn.Module):
    def __init__(self, input_dims, output_dims):
        super(ResidualBlock, self).__init__()
                # Defining the blocks
        self.conv1 = nn.Sequential( 
            nn.Conv2d(input_dims, output_dims, kernel_size = (3,3), padding = (1,1)),
            nn.ReLU()
        )
        self.conv2 = nn.Conv2d(output_dims, output_dims, kernel_size =(1,1))
        self.act = nn.ReLU()

    def forward(self, X): 
                # Clone X as this will be passed in a later layer
        X2 = X.clone()
        X = self.conv1(X)
        X = self.conv2(X)
                # Add the output of the previous layer to X2
        X = self.act(X + X2)
        return X

In the above code, we can see that we have separate attributes for the different layers and we have customized our forward function a bit. Since we needed to pass the initial input X onto a future layer, we stored its value in X2. We then transform X in place by passing it to the layers and then finally adding X2 before the final activation.

Question: How come this is possible? Don’t convolutions downsample the image thus making X + X2 not possible due to their difference in shape.

Great question! I actually asked myself this while reviewing ResNets. It turns out that the reason why this works is as simple as that these tensors have the same exact shape. The first convolution in the block particularly has a 1x1 padding and has a kernel size of 3x3, hence the output shape would be the same as the input shape. Recall: $\frac{\text{dim}+2p-k}{s}$ where dim is the input dimension, p is the padding, k is the kernel-size, and s is the stride in the specific dimension.

$\frac{\text{dim}+2p-k}{s}$

class ResNet(nn.Module):
    def __init__(self):
        super(ResNet, self).__init__() 
        self.model = nn.Sequential(
            nn.Conv2d(3, 64, stride = (2,2), kernel_size=(7,7)),  
            nn.MaxPool2d(stride = (2,2), kernel_size = (7,7)),
            ResidualBlock(64,64), 
            ResidualBlock(64,64), 
            ResidualBlock(64,64), 
            ResidualBlock(64,64), 
            ResidualBlock(64,64), 
            ResidualBlock(64,64), 
            nn.AdaptiveAvgPool2d(output_size=(7, 7)),
            classifier_block(64*7*7, 4096, 1000)

        )
    def forward(self,X):
        return self.model(X)

ResNet(
(model): Sequential(
(0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2))
(1): MaxPool2d(kernel_size=(7, 7), stride=(2, 2), padding=0, dilation=1, ceil_mode=False)
(2): ResidualBlock(
(conv1): Sequential(
(0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU()
)
(conv2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
(act): ReLU()
)
(3): ResidualBlock(
(conv1): Sequential(
(0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU()
)
(conv2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
(act): ReLU()
)
(4): ResidualBlock(
(conv1): Sequential(
(0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU()
)
(conv2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
(act): ReLU()
)
(5): ResidualBlock(
(conv1): Sequential(
(0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU()
)
(conv2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
(act): ReLU()
)
(6): ResidualBlock(
(conv1): Sequential(
(0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU()
)
(conv2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
(act): ReLU()
)
(7): ResidualBlock(
(conv1): Sequential(
(0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU()
)
(conv2): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
(act): ReLU()
)
(8): AdaptiveAvgPool2d(output_size=(7, 7))
(9): Sequential(
(0): Flatten()
(1): Linear(in_features=3136, out_features=4096, bias=True)
(2): ReLU()
(3): Dropout(p=0.5, inplace=False)
(4): Linear(in_features=4096, out_features=4096, bias=True)
(5): ReLU()
(6): Dropout(p=0.5, inplace=False)
(7): Linear(in_features=4096, out_features=1000, bias=True)
(8): Softmax(dim=None)
)
)
)

A Slightly More Complex Model

Recurrent Neural Networks are neural networks that perform better with serial data. To understand RNNs, it is integral that we understand the RNN cell first.

The RNN Cell

The RNN Cell is simply a compact linear network that accepts two inputs. It is governed by the following formulae.

$a^{} = f(W_{aa}a^{} + W_{ax}x^{} + b_a)$

$\hat y^{} = g(W_{ya}a^{} + b_y)$

Figure 3. An RNN Cell. Image retrieved from deeplearning.ai Sequence Models course.

It takes in two inputs— the activations from the previous cell, $a^{}$ and the output from the previous layer $x^{}$. These are concatenated and applied to a linear model parameterized by $W_a$ and $b_a$ to get the current cell’s activation. This activation is passed to the next cell and is also transformed linearly parameterized by $W_y$ and $b_y$(optional) to form $\hat y$.

class RNNCell(nn.Module):
    def __init__(self, embed_length,act_dim, output_dim, **kwargs):
        super(RNNCell, self).__init__()

        act = kwargs.get("act", "relu")
        acty = kwargs.get("acty", "relu")
        activation_map = {"relu": nn.ReLU(), "l-relu": nn.LeakyReLU(), "sig": nn.Sigmoid(), "tanh": nn.Tanh()}
        self.act_dim = act_dim

        self.linear = nn.Linear(act_dim+embed_length,act_dim) 
        self.activation = activation_map.get(act)
        self.linear_y = nn.Linear(act_dim, output_dim)
        self.activation_y = activation_map.get(acty)
        pass

    def forward(self, a, X): 
        # X is n x embed_length
        # a is n x act_dim 
        assert(a.shape[1] == self.act_dim)
        # concatenate a and X since they will be transformed by the same Linear.
        X = torch.cat((a,X), axis = 1)
        # transform the inputs
        X = self.linear(X)
        a = self.activation(X)
        X = self.linear_y(a)
        Y = self.activation_y(X)
        return a,Y

        pass

An RNN Layer

An RNN layer is just a series of RNN cells with its activations fed to the next unit. It’s forward function is more complicated as the previous unit’s activation is passed to the next one.

Figure 4. An RNN Layer. Image retrieved from colah’s blog.

The forward function calculates each cell’s activations and outputs first. Calculating each cell from left to right since the next cell requires the previous cell’s activations.

class RNN(nn.Module):
    def __init__(self, **kwargs):
        super(RNN, self).__init__()
        # output shape is 3d. n x output_dim x embed_length
        # input_dim is the embedding length(ie. word embedding length)
        self.input_dim = kwargs.get("input_dim", 0)
        self.act_dim = kwargs.get("act_dim", self.input_dim)
        self.output_dim = kwargs.get("output_dim", 1)
        self.time_steps = kwargs.get("time_steps", self.output_dim) 
        self.unit_output_dim = kwargs.get("unit_output_dim" , 1)

        assert(self.output_dim <= self.time_steps)
        # Populate the layer with the cells based on timesteps.
        self.models = nn.ModuleList([
            RNNCell(self.input_dim,self.act_dim, self.unit_output_dim) 
        ] * self.time_steps )

    def forward(self,  X):
        # x is n x time_steps x embed_length 
        n = X.shape[0] 

        # make sure X axis 1 is less than time_steps, cause it might crash.
        assert(X.shape[1] <= self.time_steps)

        # Sometimes our input would have lesser dimensions than our timesteps, hence we add zero tensors to it
        # to match the size and be compatible.
        X = torch.cat((X, torch.zeros(n, self.time_steps - X.shape[1], self.input_dim)), axis = 1)

        # Initialize the first activation.
        a = torch.zeros(n, self.act_dim)

        # Create the Y array, the individual y-predictions will be stored here
        Y = torch.zeros(n, self.output_dim, self.unit_output_dim)

        # Create iterator for y_i for locating which timestep we are in.
        y_i = 0
        for i,cell in enumerate(self.models):
            # Get the input for the current timestep
            x = X[:, i, :] 
            # Forward the x and activations to the current cell.
            a,y = cell(a, x)  
            # Only add the last predictions for the output of size output_dim
            if i >= self.time_steps - self.output_dim:
                Y[:, y_i,:] = y 
                y_i+=1

        return Y

A Basic RNN

That was hectic! Now let’s put it all together, a Sequential object would do for now to avoid complexity. Let’s stack these RNN Layers together.

model = nn.Sequential(

    RNN(input_dim = embed_length, act_dim = 10, time_steps = 5, output_dim = 4, unit_output_dim = embed_length),
    RNN(input_dim = embed_length, act_dim = 7, time_steps = 4, output_dim = 2, unit_output_dim = embed_length),
    RNN(input_dim = embed_length, act_dim = 3, time_steps = 2, output_dim = 1, unit_output_dim = 1)

)

Sequential(
(0): RNN(
(models): ModuleList(
(0): RNNCell(
(linear): Linear(in_features=22, out_features=10, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=10, out_features=12, bias=True)
(activation_y): ReLU()
)
(1): RNNCell(
(linear): Linear(in_features=22, out_features=10, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=10, out_features=12, bias=True)
(activation_y): ReLU()
)
(2): RNNCell(
(linear): Linear(in_features=22, out_features=10, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=10, out_features=12, bias=True)
(activation_y): ReLU()
)
(3): RNNCell(
(linear): Linear(in_features=22, out_features=10, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=10, out_features=12, bias=True)
(activation_y): ReLU()
)
(4): RNNCell(
(linear): Linear(in_features=22, out_features=10, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=10, out_features=12, bias=True)
(activation_y): ReLU()
)
)
)
(1): RNN(
(models): ModuleList(
(0): RNNCell(
(linear): Linear(in_features=19, out_features=7, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=7, out_features=12, bias=True)
(activation_y): ReLU()
)
(1): RNNCell(
(linear): Linear(in_features=19, out_features=7, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=7, out_features=12, bias=True)
(activation_y): ReLU()
)
(2): RNNCell(
(linear): Linear(in_features=19, out_features=7, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=7, out_features=12, bias=True)
(activation_y): ReLU()
)
(3): RNNCell(
(linear): Linear(in_features=19, out_features=7, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=7, out_features=12, bias=True)
(activation_y): ReLU()
)
)
)
(2): RNN(
(models): ModuleList(
(0): RNNCell(
(linear): Linear(in_features=15, out_features=3, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=3, out_features=1, bias=True)
(activation_y): ReLU()
)
(1): RNNCell(
(linear): Linear(in_features=15, out_features=3, bias=True)
(activation): ReLU()
(linear_y): Linear(in_features=3, out_features=1, bias=True)
(activation_y): ReLU()
)
)
)
)

That's it for now! We could go on here with all the different types of RNNs but this post is getting too long already. I could continue with the more complicated models if this post gains enough traction so make sure to share with your friends!
'til next time!

This article is also published in Medium.