Deploying Python 3 PyTorch Networks in Python 2 for ROS

By Thomas Weng on December 11, 2022

Scenario: you train a neural network in Python 3 with PyTorch, and want to deploy your network on a real robot using ROS. However, ROS only supports Python 2 packages, and getting ROS to work with Python 3 is difficult.1

While ROS 2 does have support for Python 3, it is still in development. In the meantime, here is a possible workaround for running Python 3 networks on ROS 1.

The workaround is to load the weights of a Python 3 PyTorch model, process the weights into a format compatible with Python 2 PyTorch, then save the output as a text file to be loaded in Python 2.

In the following gist, module_save_util.py processes the Python 3 model and saves the weights for Python2: link to gist. module_load_util.py can then be called in Python 2 to load the processed weights.

In our case, we trained networks with PyTorch 1.8.0/1.8.1 and deployed with PyTorch 1.3.1. I cannot give any guarantees on whether the performance is the same after porting the weights to Python 2, though in my own experiments I did not see a drop. That said, it has been a while since I have done this solution so your mileage may vary.

Others at the Robotics Institute have proposed the following solutions, though I haven’t tried them myself:

Credit goes to Sujay Bajracharya for developing the solution in this post! If there is code in the gist above from other sources, I am happy to give credit as well.

Full gist:


Footnotes

Tags: how-to, robotics