Reducing PyTorch dependency size and using PyTorch Models for Flask apps

While participating in a hackathon I was stuck with Deploying a Flask app that used PyTorch as one of the dependencies. Reason was the size of PyTorch. Pip installing it costed above 700 MB which was unacceptable given Heroku allows total app size of only 512 MB in the free tier. I had to figure out a way to reduce PyTorch size, but even the older versions were larger than 500MB. I have explained how to use PyTorch with reduced size along with using Models trained on it in this blog.

Coming directly to the solution. Use PyTorch CPU only versions. Does this answer the whole question? NO. If it would have I won’t be writing a blog on it right?.

How to use it precisely?

To use CPU only PyTorch as one of the requirements in your app, add these 2 lines in the requirements.txt.

-f https://download.pytorch.org/whl/torch_stable.html                                                                                                                                                             torch==1.8.0+cpu

Then pip install requirements before you start training or testing the model.

pip install -r requirements.txt

How to save the model ?

Training is pretty much same by providing x as argument for the model object. But after that you must save the model as a .pth file like this.

torch.save(model.state_dict(), "PATH/MODEL-NAME.pth")

Here model.sate_dict() saves the model weights instead of the whole model architecture. Advantage is that this method helps us include PyTorch and model in apps that can’t afford huge size while deploying and the disadvantage is that you will have to carry the model architecture to the flask app or wherever you wish to use the model. Let’s see how it works.

How to use the model ?

You must have the Model/Network class and load the model as follows

# Network Class example
class Network(nn.Module):
    def __init__(self):
        super(Network,self).__init__()
        # CNNs for rgb images 
        self.conv1= nn.Conv2d(in_channels=3,out_channels=6,kernel_size=5)
        self.conv2= nn.Conv2d(in_channels=6,out_channels=12,kernel_size=5)
        .
        .
        .
        
    def forward(self,t):
        t = t
        
        t=self.conv1(t)
        t=F.relu(t)
        t=F.max_pool2d(t,kernel_size = 2, stride = 2)
        .
        .
        .
        
        return t

Loading the saved model in Flask app

# Create an object from the class above
model = Network() 
# Loading the saved model state (.pth)
model.load_state_dict(torch.load("path/model_name.pth"))
model.eval()
y_pred = model(x)

I hope the blog has helped you to understand how to use PyTorch in a way that it doesn’t take up too much space in your dependencies.