Transpiling a Tensorflow model to build on top#
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 noise = tf.random.normal(shape=(1, 224, 224, 3)) torch_eff_encoder = ivy.transpile(eff_encoder, to="torch", args=(noise,))
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))
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])
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 ➡️