Pytorch vs Tensorflow in 2020. How the two popular frameworks have... | by Moiz Saifee | Towards Data Science

Pytorch and Tensorflow are by far two of the most popular frameworks for Deep Learning. It’s always a lot of work to learn and be comfortable with a new framework, so a lot of people face the dilemma of which one to choose out of the two
The two frameworks had a lot of major differences in terms of design, paradigm, syntax etc till some time back, but they have since evolved a lot, both have picked up good features from each other and are no longer that different.
A lot of online articles comparing the two are a little old, and do not appropriately capture the present scenario. In fact Soumith Chintala, one of the original authors of PyTorch, also recently tweeted about how the two frameworks are pretty similar now.
Tensorflow is from Google and was released in 2015, and PyTorch was released by Facebook in 2017.
Tensorflow arrived earlier at the scene, so it had a head start in terms of number of users, adoption etc but Pytorch has bridged the gap significantly over the years
Both work on fundamental data type called Tensors which are nothing but multi-dimensional arrays, amenable to high performance computation. Both represent computation as a directed acyclic graph often called Computation Graph.
Let’s look into some of the important aspect about these frameworks, the major differences in the beginning and where things stand as of today.
Programming API
The Tensorflow API was very cryptic to start with, it almost felt like learning a new programming language, on top of it, it was also very hard to debug due to its static computation graph approach (more on that below). Pytorch (python) API on the other hand is very Pythonic from the start and felt just like writing native Python code and very easy to debug.
Tensorflow did a major cleanup of its API with Tensorflow 2.0, and integrated the high level programming API Keras in the main API itself. A combination of these two significantly reduced the cognitive load which one had to undergo while writing Tensorflow code in the past :-)
The programming APIs (of TensorFlow and PyTorch) in fact look very similar now, so much that the two are indistinguishable a number of times (see example towards the end)
Computation graph was a major design difference between the two frameworks to start with.
Tensorflow adopted a static Computation Graph approach where one defines the sequence of computation that one wants to do, with placeholder for the data. After that for training / running the model you feed in the data. Static computation graph is great for performance and ability to run on different devices (cpu / gpu / tpu), but is a major pain to debug.
Pytorch on the other hand adopted a dynamic computation graph approach, where computations are done line by line as the code is interpreted. This makes it a lot easier to debug the code, and also offers other benefits — example supporting variable length inputs in models like RNN.
Fast forward to today, Tensorflow introduced the facility to build dynamic computation graph through its “Eager” mode, and PyTorch allows building of static computational graph, so you kind of have both static/dynamic modes in both the frameworks now.
Both the frameworks provided the facility to run on single / multiple / distributed CPUs or GPUs. In earlier days it used to be a pain to get Tensorflow to work on multiple GPUs as one had to manually code and fine tune performance across multiple devices, things have changed since then and now its almost effortless to do distributed computing with both the frameworks.
Google made its custom hardware accelerator Tensor Processing Unit (TPU) which can run computation at blazing speed, even a lot faster than GPU, available for 3rd party use in 2018. Since Tensorflow and TPU are both from Google, its far easier to run code on TPUs using Tensorflow as opposed to PyTorch, as PyTorch has a bit of patchy way of working on TPUs using third party libraries like XLA
Deployment is something where Tensorflow had a lot of advantage over PyTorch, in part due to better performance due to its Static Computation graph approach, but also due to packages / tools that facilitated quick deployment over cloud, browser or mobile. That’s the reason a lot of companies preferred Tensorflow when it came to production.
PyTorch has tried to bridge this gap in version 1.5+ with TorchServe, but its yet to mature
Its amusing that for a lot of things the APIs are so similar that the codes are almost indistinguishable. Given below are code snippets for the core components on MNIST Digit Recognition (proverbial “Hello World” problem in Computer Vision) for both Tensorflow and Pytorch, try to guess which one is which
The complete Tensorflow and Pytorch code is available at my Github Repo
Data Loader
# Download the MNIST Data(x_train, y_train), (x_test, y_test) = mnist.load_data()x_train, x_test = x_train / 255.0, x_test / 255.0# Add a channels dimensionx_train = x_train[..., tf.newaxis]x_test = x_test[..., tf.newaxis]train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000).batch(32)test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)
Which one is PyTorch code - above or below?
# Download the MNIST Data and create dataloadertransform = transforms.Compose([transforms.ToTensor()])xy_train = datasets.MNIST('./', download=True, train=True, transform=transform)xy_test = datasets.MNIST('./', download=True, train=False, transform=transform)train_ds = DataLoader(xy_train, batch_size=32, shuffle=True)test_ds = DataLoader(xy_test, batch_size=32, shuffle=True)
Okay the method to load the data looks a bit different, but I promise it gets similar from here on :-)
Model Definition
# Model Definitionclass MyModel(nn.Module):    def __init__(self):        super(MyModel, self).__init__()        self.conv1 = Conv2d(in_channels=1, out_channels=32, kernel_size=3)        self.flatten = Flatten()        self.d1 = Linear(21632, 128)        self.d2 = Linear(128, 10)    def forward(self, x):        x = F.relu(self.conv1(x))        x = self.flatten(x)        x = F.relu(self.d1(x))        x = self.d2(x)        output = F.log_softmax(x, dim=1)        return output
Which one is PyTorch code - above or below?
# Model Definitionclass MyModel(Model):    def __init__(self):        super(MyModel, self).__init__()        self.conv1 = Conv2D(filters=32, kernel_size=3, activation='relu')        self.flatten = Flatten()        self.d1 = Dense(128, activation='relu')        self.d2 = Dense(10)    def call(self, x):        x = self.conv1(x)        x = self.flatten(x)        x = self.d1(x)        output = self.d2(x)        return output
Instantiate Model, Loss, Optimizer
# Instantiate Model, Optimizer, Lossmodel = MyModel()optimizer = Adam()loss_object = SparseCategoricalCrossentropy(from_logits=True, reduction='sum')
Which one is PyTorch code - above or below?
# Instantiate Model, Optimizer, Lossmodel = MyModel()optimizer = Adam(model.parameters())loss_object = CrossEntropyLoss(reduction='sum')
Training Loop
for epoch in range(2):    # Reset the metrics at the start of the next epoch    train_loss = 0    train_n = 0    for images, labels in train_ds:        with GradientTape() as tape:            # training=True is only needed if there are layers with different            # behavior during training versus inference (e.g. Dropout).            predictions = model(images, training=True)            loss = loss_object(labels, predictions)            train_loss += loss.numpy()            train_n += labels.shape[0]        gradients = tape.gradient(loss, model.trainable_variables)        optimizer.apply_gradients(zip(gradients, model.trainable_variables))    train_loss /= train_n
Which one is PyTorch code - above or below?
for epoch in range(2):    # Train    model.train()    train_loss = 0    train_n = 0    for image, labels in train_ds:        predictions = model(image).squeeze()        loss = loss_object(predictions, labels)        train_loss += loss.item()        train_n += labels.shape[0]        loss.backward()        optimizer.step()        optimizer.zero_grad()    train_loss /= train_n
It’s really interesting (and convenient!) how similar the APIs seem to be now.
Tensorflow and PyTorch are two excellent frameworks for research and development of deep learning applications. They do the heavy lifting in terms of computation, managing the underlying hardware and have huge communities which makes it a lot easier to develop custom application by standing on the shoulder of giants.
There were some major differences in the two frameworks till a while back, both have since adopted good features from each other and have both become better in the process. There are still things which are slightly easier in one compared to another, but its now also easier than ever to switch back and forth between the two due to increased similarity.