PyTorch is too heavy for deployment, is there a light weight pytorch for running neural networks?

Question

We deploy PyTorch models in docker container, which massively increased the size of the docker container by more than 1G.

But when we deploy the model the training has already been done, so technically we don't need to include the machinery involved in training.

We don't even need backpropagation, we just need to run the neural network to get the outputs.

Is it possible to include only part of PyTorch (or another totally different product) that executes a neural network and nothing else? So deployment is light-weight.

Answer 1

One thing to try is to create a lean version of Pytorch that only has what you need. I did something similar for an AWS Lambda layer, here's what I was able to delete (this was for Pytorch 1.1, things may have changed since):

find . -type d -name "tests" -exec rm -rf {} +
find . -type d -name "__pycache__" -exec rm -rf {} +
rm -rf ./{caffe2,wheel,pkg_resources,boto*,aws*,pip,pipenv,setuptools} 
rm ./torch/lib/libtorch.so 
rm -rf ./torch/test

If that's not enough, you need to look into using ONNX or TorchScript/torch.jit. ONNX exports your model to a binary protobuf file. TorchScript exports it to a serialized form that can be run in a C++ environment.