Step-by-Step Procedure of Converting TensorFlow Model to PyTorch Model
Setting Up the Environment
Let’s make sure everything is configured properly in our environment before beginning the conversion procedure. Install the required packages by using:
!pip install tensorflow torch
Create a TensorFlow Model
Python3
import numpy as np from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split from sklearn.preprocessing import OneHotEncoder import tensorflow as tf from tensorflow.keras.models import Sequential from tensorflow.keras.layers import Dense # Load the Iris dataset iris = load_iris() X = iris.data y = iris.target.reshape( - 1 , 1 ) # Reshape to make it a column vector # One-hot encode the target variable encoder = OneHotEncoder(categories = 'auto' ) y = encoder.fit_transform(y).toarray() # Split the dataset into training and testing sets X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2 , random_state = 42 ) # Step 1: Define the model model = Sequential([ Dense( 10 , activation = 'relu' , input_shape = (X_train.shape[ 1 ],)), Dense( 8 , activation = 'relu' ), Dense( 3 , activation = 'softmax' ) ]) model.summary() |
Output:
Model: "sequential_3"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
dense_8 (Dense) (None, 10) 50
dense_9 (Dense) (None, 8) 88
dense_10 (Dense) (None, 3) 27
=================================================================
Total params: 165 (660.00 Byte)
Trainable params: 165 (660.00 Byte)
Non-trainable params: 0 (0.00 Byte)
Train and Save the Model
Python3
#Compile the model model. compile (optimizer = 'adam' , loss = 'categorical_crossentropy' , metrics = [ 'accuracy' ]) #Train the model model.fit(X_train, y_train, epochs = 100 , batch_size = 4 , verbose = 1 ) #Evaluate the model loss, accuracy = model.evaluate(X_test, y_test) print (f 'Test Loss: {loss:.4f}' ) print (f 'Test Accuracy: {accuracy:.4f}' ) #Save the model model.save( 'iris_model.h5' ) |
Output:
Epoch 1/100
30/30 [==============================] - 2s 2ms/step - loss: 1.1517 - accuracy: 0.3417
Epoch 2/100
30/30 [==============================] - 0s 2ms/step - loss: 1.0865 - accuracy: 0.4000
Epoch 3/100
30/30 [==============================] - 0s 2ms/step - loss: 1.0580 - accuracy: 0.4833
Epoch 4/100
30/30 [==============================] - 0s 2ms/step - loss: 1.0397 - accuracy: 0.4500
Epoch 5/100
30/30 [==============================] - 0s 2ms/step - loss: 1.0172 - accuracy: 0.3917
..
Test Loss: 0.0591
Test Accuracy: 1.0000
Load the trained TensorFlow model
Python3
loaded_model = tf.keras.models.load_model( "iris_model.h5" ) |
Converting to PyTorch Model
Installing the Required Libraries
In order to convert TensorFlow models to ONNX format, install the tf2onnx library:
!pip install tf2onnx
!pip install onnx2pytorch
Converting to tf2onnx Model
Python3
import tf2onnx # Convert the model to ONNX format onnx_model, _ = tf2onnx.convert.from_keras(loaded_model) |
Converting to PyTorch Model
Python3
import onnx from onnx2pytorch import ConvertModel # Convert ONNX model to PyTorch pytorch_model = ConvertModel(onnx_model) pytorch_model |
Output:
ConvertModel(
(MatMul_sequential_2/dense_5/BiasAdd:0): Linear(in_features=4, out_features=10, bias=True)
(Relu_sequential_2/dense_5/Relu:0): ReLU(inplace=True)
(MatMul_sequential_2/dense_6/BiasAdd:0): Linear(in_features=10, out_features=8, bias=True)
(Relu_sequential_2/dense_6/Relu:0): ReLU(inplace=True)
(MatMul_sequential_2/dense_7/BiasAdd:0): Linear(in_features=8, out_features=3, bias=True)
(Softmax_dense_7): Softmax(dim=-1)
)
How to Convert a TensorFlow Model to PyTorch?
The landscape of deep learning is rapidly evolving. While TensorFlow and PyTorch stand as two of the most prominent frameworks, each boasts its unique advantages and ecosystems.
However, transitioning between these frameworks can be daunting, often requiring tedious reimplementation and adaptation of models. Fortunately, the Open Neural Network Exchange (ONNX) format emerges as a powerful intermediary, facilitating smooth conversions between TensorFlow and PyTorch models.
In this article, we will learn how can we use ONNX to convert TensorFlow model into a Pytorch model.