A tale of two frameworks: PyTorch vs. TensorFlow
Comparing auto-diff and dynamic model sub-classing approaches with PyTorch 1.x and TensorFlow 2.x
The data science community is a vibrant and collaborative space. We learn from each other’s publications, debate ideas on forums and online outlets, and share lots (and lots) of code. A natural side effect of this collaborative spirit is the high likelihood of encountering unfamiliar tools used by our colleagues. Because we don’t work in a vacuum, it often makes sense to gain familiarity with several languages and libraries in a given subject area in order to collaborate and learn most effectively.
It’s no surprise, then, that many data scientists and Machine Learning engineers have two popular Machine Learning frameworks in their toolboxes: TensorFlow and PyTorch. These frameworks — both in Python — share many similarities, but also diverge in meaningful ways. These differences, such as how they handle APIs, load data, and support specialized domains, can make alternating between the two frameworks cumbersome and inefficient. And that’s a problem given how common both of these tools are.
Therefore, this article aims to illustrate the differences between PyTorch and TensorFlow by focusing on the basics of creating and training two simple models. In particular, we’ll cover how to use dynamic subclassed models with the Module API from PyTorch 1.x and the Module API from TensorFlow 2.x. We’ll look at how auto-diff can be used by these frameworks, too, to provide very naïve implementations of gradient descent.
But first, data
Because we are focusing on the core of auto-diff/auto-grad functionality (which, as a refresher, is the capacity to automatically extract the derivative of a function and apply the gradients with respect of some parameters in order to optimize these parameters using gradient descent) we can start with the simplest possible model, a linear regression. We can use the numpy library to generate some linear data with a bit of random noise, and then run our models on that dummy data set.
The models
Once we have data, we can implement a regression model from scratch in both TensorFlow and PyTorch. For simplicity, we won’t initially use any layers or activators, just define two tensors, w and b, that represent the weights and bias of the linear model y = wx + b.
As you can see below, other than a few differences in the API names, the class definitions for the two models are nearly identical. The most significant difference is that PyTorch requires an explicit Parameter object to define the weights and bias tensors to be captured by the graph, whereas TensorFlow is able to automagically capture those same parameters. Indeed, PyTorch parameters are tensor subclasses that have a special property when used with the Module API: They automatically add self to the list of Module parameters, so self appears, for example, in the parameters() iterator.
Both frameworks extract everything needed to generate the graph from this class definition and its execution methods (__call__ or forward) and, as we see below, calculate the gradients needed to implement backpropagation.
TensorFlow Dynamic Model
PyTorch Dynamic Model
Building the training loop, backpropagation, and optimizers
With these simple TensorFlow and PyTorch models established, the next step is to implement the loss function, which in this case is just means squared error. We can then instantiate the model classes and run the training loop for a few cycles.
Again, because we’re focusing on the core auto-diff/auto-grad functionality, the aim here is to build a custom training loop using TensorFlow and PyTorch–specific auto-diff implementations. These implementations calculate gradients for the simple linear function and manually optimize the weight and bias parameters with a naïve gradient descent optimizer, essentially minimizing the loss computed between the real point and the prediction using the differentiable function at each point.
For the TensorFlow training loop, I explicitly used the GradientTape API to track the forward execution of the model and step-wise loss calculations. I used the gradients from GradientTape to optimize the weights and bias parameters. PyTorch provides a more “magical” auto-grad approach, implicitly capturing any operations on the parameter tensors and providing the gradients to use for optimizing the weights and bias parameters without having to call another API. Once I have the weights and bias gradients, implementing the custom gradient descent method on both PyTorch and TensorFlow is as simple as subtracting the weight and bias parameters from these gradients, multiplied by a constant learning rate.
Note that because PyTorch implements auto-diff/auto-grad automatically, it’s necessary to explicitly call the no_grad API right after we computed the backward propagation. This instructs PyTorch not to compute the gradients for the update operation of the weights and bias parameters. We also need to explicitly zero-out the previously automatically-computed gradients, computed in the forward operation, to stop PyTorch from automatically accumulating gradients in all batch and cycles iterations.
TensorFlow training loop
PyTorch training loop
PyTorch and TensorFlow models reusing available Layers
Now that I’ve shown how to implement linear regression models from scratch in PyTorch and TensorFlow, we can look at how to reimplement the same models using Dense and Linear layers, respectively, from the TensorFlow and PyTorch libraries.
TensorFlow and PyTorch dynamic models with existing layers
You’ll notice in both model initialization methods that we are replacing the explicit declaration of the w and b parameters with a Dense layer in TensorFlow and a Linear layer in PyTorch. Both of these layers implement the linear regression, and we’ll instruct them to use single weight and bias parameters in place of the explicit w and b parameters used before. Dense and Linear implementations will internally use the same tensor declarations we used before (respectively tf.Variable and nn.Parameter) to allocate these tensors and associate them with the model parameters list.
We’ll also update the call / forward methods of these new model classes to replace the manual linear regression computation with the execution of the Dense / Linear layers.
Training with available optimizers and loss functions
Now that we have re-implemented our TensorFlow and PyTorch models using existing Layers, we can focus on how to build more optimized training loops. Rather than using our previous naïve implementation, we’ll use the native Optimizers and Loss Functions available from these libraries.
We’ll continue to use the auto-diff/auto-grad functionalities observed before, but this time with a Standard Gradient Descent (SGD) optimization implementation alongside a standard loss function.
TensorFlow training loop with easy Fit method
In TensorFlow, fit() is a very powerful and high-level method for training a model. It allows us to replace a manual training loop with a single method that specifies the hyper-tuning parameters. Before calling fit() we’ll compile our model class using the compile() method, and then pass a gradient descendent optimizer and a loss function to use for the training.
You’ll notice that in this case we will reuse the methods from the TensorFlow library as much as possible. In particular, we will pass a standard Stochastic Gradient Descent (SGD) optimizer and a standard Mean Absolute Error loss function implementation (mean_absolute_error) to the compile method. Once the model is done compiling, we can finally call the fit method to fully train our model. We’ll pass the data (x and y), the number of epochs, and the batch size to use on each epoch.
TensorFlow training loop with custom loop and SGD Optimizer
In the following code snippet, we will implement another custom training loop for our model, this time reusing the loss functions and optimizers provided by the TensorFlow library as much as possible. You’ll notice that our former custom Python loss function is replaced with the tf.losses.mse() method. Instead of manually updating the model parameters with the gradients, we initialize a tf.keras.optimizers.SGD() optimizer. Calling the optimizer.apply_gradient() and passing a list of weight and bias tuples will update the model parameters with gradients.
PyTorch training loop with custom loop and SGD Optimizer
As in the previous TensorFlow code snippet above, the following code snippet implements a PyTorch training loop for our new model by reusing loss functions and optimizers provided by the PyTorch library. You’ll notice we’ll replace our former custom Python loss function with the nn.MSELoss() method and initialize a standard optim.SGD() optimizer with the list of our model’s learning parameters. As previously illustrated, we will instruct PyTorch to obtain the associated gradients for each parameter tensor from the loss backward propagation (loss.backward()), and finally, we can easily update the new standard optimizer with all parameters associated with the gradients by calling the optimizer.step() method. The way that PyTorch makes automatic associations between tensors and their gradient allows the optimizer to retrieve the tensors and gradients to update them with the configured learning rate.
Results
As we saw, the TensorFlow and PyTorch auto-diff and Dynamic sub-classing APIs are very similar, even in the way they use standard SGD and MSE implementations. Naturally, both models also gave us very similar results.
In the code snippet below, we use Tensorflow’s training_variables and PyTorch’s parameters methods to get access to the models’ parameters and plot the graph of our learned linear functions.
Conclusion
Both PyTorch and the new TensorFlow 2.x support Dynamic Graphs and auto-diff core functionalities to extract gradients for all parameters used in a graph. You can easily implement a training loop in Python with any loss functions and gradient descendent optimizers. In order to focus on real core differences between the two frameworks, we simplified the example above by implementing our own simple MSE and naïve SGD.
However, I strongly suggest that you reuse optimized and specialized code available on these libraries before implementing any naïve code.
The table below summarizes all the differences we noted in the sample code above. I hope it can serve as a useful reference for when you need to switch between these two frameworks.
Source code
Want to know more about how you can set up your Machine Learning projects for success? Read our recent article:
Jacopo Mangiavacchi is on LinkedIn.
