Can a production ready Pytorch 1.0 give TensorFlow a tough time?

PyTorch has announced a preview of the blueprint for PyTorch 1.0, the next major release of the framework. This breakthrough version is expected to bring more stability, integration support and complete production backing allowing developers to move from core research to production in an amicable way without having to deal with any migration challenges.

PyTorch is an open-source Python-based scientific computing package which provides powerful GPU acceleration. PyTorch is known for advanced indexing and functions, imperative style, integration support and API simplicity. This is one of the key reasons why developers prefer PyTorch for research and hackability. To know more about how Facebook-based PyTorch competes with Google’s TensorFlow read our take on this deep learning war.

Some of the noteworthy changes in the roadmap for PyTorch 1.0 are:

Production support

One of the biggest challenges faced by developers in terms of using PyTorch is production support. There are n number of issues faced while trying to run the models efficiently in production environments.

Even though PyTorch provides excellent simplicity and flexibility, due to its tight coupling to Python, the performance at production-scale is a challenge.

To counter these challenges, the PyTorch team has decided to bring PyTorch and Caffe2 together to provide production-scale readiness to the developers.

However, adding production support brings complexity and configurable options for models in the API. The PyTorch team will stick to the goal of keeping the platform — a favorable choice — for researchers and developers. Hence, they are introducing a new just-in-time (JIT) compiler, named torch.jit.

torch.jit compiler rewrites PyTorch models during runtime in order to achieve scalability and efficiency in production environments. It can also export PyTorch models to run in a C++ environment. (runtime based on Caffe2 bits)

Note: In PyTorch version 1.0, your existing code will continue to work as-is.

Let’s go through how JIT compiler can be used to export models to a Python-less environment in order to improve their working performance.

torch.jit: The go-to compiler for your PyTorch models

Building models using Python code, no doubt gives maximum productivity and makes PyTorch very simple and easy-to-use. However, this also means PyTorch finding it difficult to know which operation you will run next.

This can be frustrating for the developers during model export and automatic performance optimizations because they need to be aware of how the computations will look like before it even gets implemented.

To deal with these issues, PyTorch provides two ways of recovering information from the Python code. Both these methods will be useful based on different contexts, giving you the leverage to use/mix them with ease.

Tracing the native Python code

Compiling a subset of the Python language

Tracing mode

torch.jit.trace function allows you to record the native PyTorch operations performed along with the data dependencies between them.

PyTorch version 0.3 already had a tracer function which is used to export models through ONNX. This new version uses a high-performance C++ runtime that allows PyTorch to re-execute programs for you.

The key advantage of using this method is that it doesn’t have to deal with how your Python code is structured since we only trace through native PyTorch operations.

Script mode

PyTorch team has come up with a solution called scripting mode made specially for those models such as RNNs which make use of control flow.

However, you will have to write out a regular Python function (avoiding complex language features)

In order to get your function compiled, you can assign @script decorator. This will make sure it alters your Python function directly into high-performance C++ during runtime.

Advantages in optimization and export techniques

Irrespective of you using a trace or a script function, the technique allows you to optimize/export the model for use in production environments (i.e. Python-free portrayal of the model)

Now you can derive bigger segments of the model into an intermediate representation to work with sophisticated models.

You can use high-performance backends available in Caffe2 to run the models efficiently

Usability

If you don’t need to export or optimize your model, you do not need to use these set of new features.

These modes will be included into the core of the PyTorch ecosystem, thus allowing you to mix and match them with the existing code seamlessly as per your needs.

Additional changes and improvements

In addition to the major update in the production support for 1.0, PyTorch team will continue working on optimizing, working on the stability of the interface, and fixing other modules in PyTorch ecosystem

PyTorch 1.0 will see some changes in the backend side which might affect user-written C and C++ extensions.

In order to incorporate new features and optimization techniques from Caffe2, PyTorch team is replacing (optimizing) the backend ATen library.

PyTorch team is planning to release 1.0 during the summer. For the detailed preview of the roadmap, you can refer the official PyTorch blog.