PyTorch to ONNX Conversion

When converting a PyTorch model to ONNX, it is important to recognize that there are two components of a PyTorch, the model and the weights. Most often the weights are saved in a .pkl or .pth file and these weights are loaded into a model generated in PyTorch. To export a model to ONNX both of these components are necessary, so it will not be possible to generate an ONNX model with only a .pth or .pkl file. The conversion is handled by the PyTorch function https://pytorch.org/docs/stable/onnx.html#torch.onnx.export.

Below is a code snippet of an entire conversion, which will be broken down and discussed further in the article.

model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], bins)

saved_state_dict = torch.load(weights_file, map_location=torch.device('cpu'))
model.load_state_dict(saved_state_dict)

input_shape = [int(v) for v in input_shape.split(',')]
input_val = torch.randn(*input_shape)
torch.onnx.export(model, input_val, output_filename, opset_version=11)

The first section is the generation of the model object which contains the architecture of the model and where the weights will be loaded. In the code snippet, this occurs here

model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], bins)

where we generate a HopeNet model architecture, but any model architecture can be created. For further information on PyTorch model generation, please see here.

The second section of the code snippet

saved_state_dict = torch.load(weights_file, map_location=torch.device('cpu'))
model.load_state_dict(saved_state_dict)

loads the provided weights file into the model. The map_location does not have to be 'cpu', but as we are converting this to ONNX and not running inference on the PyTorch model, the map_location is unimportant and 'cpu' guarantees support.

The final preparatory section is setting the expected input value.

input_shape = [int(v) for v in input_shape.split(',')]
input_val = torch.randn(*input_shape)

The majority of PyTorch models do not have a set input resolution, so this value can be set to whatever is wanted by the user, however some models work better and are trained for certain resolutions, so it is recommended to find the input shape that fits with the model and the training performed on the weights. You do need to initialize a PyTorch tensor with the set shape, which is done with the torch.randn function.

Finally, we will export the model

torch.onnx.export(model, input_val, output_filename, opset_version=11)

providing the model with the loaded weights, the input tensor and the filename of the output ONNX model. The last parameter that is recommended to be paid close attention to is the opset_version parameter as if you are looking to quantize the ONNX model, it is best to set the version to at least 11. There are also some operations that are not supported in much later opset versions, so for conversion to RTM, it is recommended to use either 11 or 13. Other opset versions will work as well, but you may run into issues at other steps of the pipeline. There are also plenty of other parameters that can be provided, for use of those, please read the documentation on the torch.onnx.export function, linked at the top.