TensorFlow vs PyTorch — Convolutional Neural Networks (CNN) | by Gurucharan M K | Towards Data Science

In my previous article, I had given the implementation of a Simple Linear Regression in both TensorFlow and PyTorch frameworks and compared their results. In this article, we shall go through the application of a Convolutional Neural Network (CNN) on a very famous Fashion MNIST dataset using both the frameworks and compare the results.
Let us get a brief idea on both the frameworks and their history. Firstly, PyTorch is an open source machine learning library based on the Torch library. PyTorch was primarily developed by Facebook’s AI Research lab (FAIR). It is free and open-source software.
On the other hand, TensorFlow was developed by the Google Brain team for internal Google research purpose. It is widely used for machine learning applications such as neural networks. It is also free and open-source software.
In order to see which framework is more efficient and simpler to use, we shall build a ConvNet using both the frameworks and compare them. For the TensorFlow, we shall be using the Keras library.
For comparison of both the frameworks, we will use the famous Fashion-MNIST dataset. It is a dataset of Zalando’s article images consisting of a training set of 60,000 examples and a test set of 10,000 examples. Each example is a 28×28 grayscale image, associated with a label from 10 classes such as shirt, bag, sandal, coat, sneaker etc. We will build a CNN network with a common architecture of LeNet with both the frameworks and compare their results.
In the first step, we will be importing the required libraries for both the TensorFlow and PyTorch frameworks.
In the next step, we will get the Fashion MNIST dataset. As it is a very common dataset, both the frameworks have preloaded functions that enable us to download the dataset and use it.
In PyTorch, the built-in module is used to load the dataset. As we need the data in the form of tensors,the data is converted into tensors by using thetorchvision.transforms.ToTensor()and assigning it to the variable transforms. The data is then downloaded with torchvision.datasets. The data is correspondingly split into the train_dataset and the test_dataset.
On the other hand, getting the data from the keras library using TensorFlow is more simpler compared to the PyTorch version. The dataset is loaded from keras.datasets and it is split into the train_images and test_images accordingly.
In this step, we shall load the data into the training set and test set. Additionally, we shall also view how the image is stored in both the frameworks.
In PyTorch, the images are stored in the ranges 0–1 and hence the images are displayed in B/W. The images are loaded into the train_loader and the test_loader using the DataLoaderfunction. Additionally, the label of the image is also printed.
On the other hand, TensorFlow stores the images in the range 0–255 and hence there is a color gradient in the same image. TensorFlow is more simple in its approach of loading the data.
In this step, we shall build the CNN Model. For this, we will build the standard LeNet-5 Architecture which is a very simple Architecture that was built in the year 1998. Given below is the architecture of LeNet-5.
In this, we build the LeNet-5 Architecture using the PyTorch framework. It has 2 Convolutional layers each followed by a Average Pooling Layer, 2 Fully connected layers and a final output classifier layer with 10 classes as the final output has 10 categories of fashion items.
In this, the model is built by creating a special class ‘NeuralNet’ where we have to define each layer of the network/model using the __init__ method. It is important to call this function ‘forward’ as this will override the base forward function in nn.Module and allow all the nn.Module functionality to work correctly.
On the other hand, building the same model in TensorFlow doesn’t require creating a class even though we use the Sequential() method. We find that building the same model with TensorFlow framework is much more simpler and explicit than building the same model in a PyTorch framework first hand.
In this step, we shall compile the model built in the previous step and set the optimizer and the loss functions. For this, in both the frameworks, we will use the Adam optimizer and the Categorical Cross Entropy loss function as there are more than 2 classes in the dataset.
In PyTotch, we first create an instance of the ‘NeuralNet’ model built in the last step as modelpy. We then assign the optimizer and the loss function to the corresponding variables. Then we visualize the model by printing the modelpy.
On the other hand, in the TensorFlow framework, we use the compile() function in which we will define the optimizer and the loss function as defined above. Then we use the summary() to see the layers of the LeNet-5 Architecture built above.
It is in this step that we deploy the models built for training on the data in the training set. For both the frameworks, we will train the model for 30 epochs and see the results.
For PyTorch, we have to individually compute all the steps including the loss function, optimizer, back-propagation and update each step in a for loop.
In TensorFlow, the Fashion MNIST is an array with two dimensions, we have to reshape our training dataset to a tensor and convert the labels from integers to binary class matrix. The model is trained using the fit() function which consists of the defined number of epochs.
In this final step, we shall compare the results of both the models built using the frameworks on the test set with the accuracy.
In PyTorch, the output prediction is predicted by outputs = modelpy(images). The accuracy of the test set is then calculated and displayed.
In TensorFlow, the predict() function is used to predict the classes of the images in the test dataset and the accuracy is calculated and displayed.
The Test Set consists of 1000 images and the models built on both the PyTorch and the TensorFlow frameworks are used to predict the classes of the test set images.
>>Test Accuracy of the model on the 10000 test images: 89% with PyTorch>>Test Accuracy of the model on the 10000 test images: 91.34% with TensorFlow
From the results, we see that the test set accuracy of the TensorFlow model is 91.34% which is a little higher than that of the PyTorch model with 89%.
From the above comparison, we were able to get an understanding of building a CNN model using two of the most popular frameworks used in Deep Learning today.
In my opinion, if you are a beginner and are totally confused as to which framework to use, I’d suggest you start with the TensorFlow framework as it is simpler and easy to understand. On the other hand, PyTorch is more fast and dynamic but needs a bit of understanding of Python API.
I have attached my GitHub repository where you can find the entire code of this comparison.
github.com
Do check out my previous article where I had given a comparison of both the frameworks in building a Simple Linear Regression model.
towardsdatascience.com
I do hope that I was able to give a good comparison of the two most popular frameworks that are used in Deep Learning today. Meet you all in my next article. Till then, Happy Machine Learning!