Transpiling Functions from PyTorch to TensorFlow

You can install the dependencies required for this notebook by running the cell below ⬇️, or check out the Get Started section of the docs to find out more about installing ivy.

!pip install ivy
!pip install torch
!pip install tensorflow
!pip install kornia

Here we’ll go through an example of how any function using torch can be converted, and used in, tensorflow via ivy.transpile. We’ll use kornia as our example, which is a state-of-the-art computer vision library built on top of torch.

First, some boiler plate imports:

import ivy
import kornia
import numpy as np
import tensorflow as tf
import torch

Now, lets transpile a kornia function that we want to use in tensorflow. The ivy.transpile call returns a new tensorflow function which is mathematically equivalent to the torch function we passed. This can take up to a minute to run.

tf_rgb_to_grayscale = ivy.transpile(kornia.color.rgb_to_grayscale, source="torch", target="tensorflow")

We can now use this function exactly as the original kornia function would be, just passing tensorflow tensors rather than torch tensors:

# using the original torch function
kornia.color.rgb_to_grayscale(torch.rand(1, 3, 28, 28))

# using the transpiled tensorflow function
tf_rgb_to_grayscale(tf.random.uniform((1, 3, 28, 28))).shape

TensorShape([1, 1, 28, 28])

Finally, lets check that the outputs of both the original torch function and transpiled tensorflow are identical when given the same inputs.

torch_x = torch.rand(1, 3, 28, 28)
tf_x = tf.convert_to_tensor(torch_x.numpy())

torch_out = kornia.color.rgb_to_grayscale(torch_x)
tf_out = tf_rgb_to_grayscale(tf_x)

np.allclose(torch_out.numpy(), tf_out.numpy(), atol=1e-6)

True