Ctrl+K

Video Tutorial


Transpiling a Tensorflow model to build on top#

Transpile a keras model to torch and build a new model around it.

⚠️ If you are running this notebook in Colab, you will have to install Ivy and some dependencies manually. You can do so by running the cell below ⬇️

If you want to run the notebook locally but don’t have Ivy installed just yet, you can check out the Get Started section of the docs.

[ ]:

%pip install ivy

In Transpile Any Model we have seen how to transpile a torch model, now let’s do the same with a model from keras transpiling a model from Keras to Torch and building a classifier on top of the resulting module.

As usual, let’s start with the imports

[ ]:

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

Now, instead of building our own Keras model, we will get one directly from Keras.

In this case, we are going to use a EfficientNet. We can download the pretrained weights with weights="imagenet" and set include_top=False to only retrieve the feature extractor.

[ ]:

# Get a pretrained keras model
eff_encoder = tf.keras.applications.efficientnet_v2.EfficientNetV2B0(
    include_top=False, weights="imagenet", input_shape=(224, 224, 3)
)

Downloading data from https://storage.googleapis.com/tensorflow/keras-applications/efficientnet_v2/efficientnetv2-b0_notop.h5
24274472/24274472 [==============================] - 0s 0us/step

Now, we will transpile the EfficientNet feature extractor to PyTorch using ivy.transpile and passing a sample tf.tensor with noise.

[ ]:

# Transpile it into a torch.nn.Module with the corresponding parameters
ivy.set_soft_device_mode(True)
noise = tf.random.normal(shape=(1, 224, 224, 3))
torch_eff_encoder = ivy.transpile(eff_encoder, source="tensorflow", to="torch", args=(noise,), modes_to_trace="eval")

To ensure that the transpilation has been correct, let’s check with a new input in both frameworks. Keep in mind that all the functions called within torch_eff_encoder are now PyTorch functions 🔀

[ ]:

x = np.random.random(size=(1, 224, 224, 3)).astype(np.float32)
output_tf = eff_encoder(tf.constant(x, dtype=tf.float32))
output_torch = torch_eff_encoder(torch.tensor(x))
print(np.allclose(output_tf , output_torch.detach().numpy(), rtol=1e-1))

True

Now, we can build or own classifier using the transpiled module as the feature extractor:

[ ]:

class Classifier(torch.nn.Module):
    def __init__(self, num_classes=20):
        super(Classifier, self).__init__()
        self.encoder = torch_eff_encoder
        self.fc = torch.nn.Linear(1280, num_classes)

    def forward(self, x):
        x = self.encoder(x)
        return self.fc(x)

And finally, we can use our new model! As we have mentioned in “Learn the Basics”, the transpiled model is fully trainable in the target framework, so you can also fine-tune your transpiled modules or train them from the ground up! 📉

[ ]:

model = Classifier()
x = torch.randn(1, 224, 224, 3)
ret = model(x)
print(type(ret), ret.shape)

<class 'torch.Tensor'> torch.Size([1, 7, 7, 20])

Round Up#

That’s it! Now you are ready to transpile any TensorFlow model, layer or trainable module and integrate it within PyTorch, but let’s keep exploring how we can convert trainable modules from (and to!) other frameworks ➡️

On this page