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

WARNING:root:   Some binaries seem to be missing in your system. This could be either because we don't have compatible binaries for your system or that newer binaries were available. In the latter case, calling ivy.utils.cleanup_and_fetch_binaries() should fetch the binaries binaries. Feel free to create an issue on https://github.com/ivy-llc/ivy.git in case of the former

WARNING:root:
Following are the supported configurations :
compiler : cp38-cp38-manylinux_2_17_x86_64, cp38-cp38-win_amd64, cp39-cp39-manylinux_2_17_x86_64, cp39-cp39-win_amd64, cp310-cp310-manylinux_2_17_x86_64, cp310-cp310-win_amd64, cp310-cp310-macosx_12_0_arm64, cp311-cp311-manylinux_2_17_x86_64, cp311-cp311-win_amd64, cp311-cp311-macosx_12_0_arm64, cp312-cp312-manylinux_2_17_x86_64, cp312-cp312-win_amd64, cp312-cp312-macosx_12_0_arm64

2024-10-18 14:51:14.248738: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-10-18 14:51:14.406321: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-10-18 14:51:14.472277: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-10-18 14:51:14.492596: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-10-18 14:51:14.630290: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-10-18 14:51:15.705472: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT

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")

Transpiling rgb_to_grayscale from torch to tensorflow. This could take a few.

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1729245116.107734  357721 cuda_executor.cc:1001] could not open file to read NUMA node: /sys/bus/pci/devices/0000:01:00.0/numa_node
Your kernel may have been built without NUMA support.
2024-10-18 14:51:56.112350: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2343] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...

Transpilation of rgb_to_grayscale complete.

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