Prepare Custom Model
In this tutorial, we will learn how to prepare a custom PyTorch model to integrate it in a PlayTorch demo.
This section will guide you step-by-step for how to export a ScriptModule with model weights for the PyTorch Mobile Lite Interpreter runtime, which is used by PlayTorch to run inference with ML models. As an example, we will export the MobileNet V3 (small) model.
Set up Python virtual environment
It is recommended to run the Python scripts in a virtual environment. Python offers a command to create a virtual environment with the following command.
python3 -m venv venv
The Python script requires
torchvision. Use the Python package manager (
pip3 or simply
pip in a virtual environment) to install both dependencies.
pip3 install torch torchvision
Export the MobileNet V3 Model
The following script will download the MobileNet V3 model from the PyTorch Hub, optimize it for mobile use, and save it in a
ptl format file (for PyTorch Mobile Lite Interpreter runtime).
from torch.utils.mobile_optimizer import optimize_for_mobile
model = torchvision.models.mobilenet_v3_small(pretrained=True)
scripted_model = torch.jit.script(model)
optimized_model = optimize_for_mobile(scripted_model)
print("model successfully exported")
export_model.py file, add the Python script above, and then run Python script to create
The script will create a file
mobilenet_v3_small.ptl. Upload the
mobilenet_v3_small.ptl file to a server that's publicly accessible, and then head over to the image classification tutorial to see how you can use it with PlayTorch.