Introduction

PyTorch and TensorFlow solve the same deep-learning problems with different APIs and slightly different defaults. The easiest way to compare them is feature by feature: tensors, autograd, model definition, data input, training loops, saving, and device placement.

Tensor Creation and Basic Ops

The lowest-level concepts map fairly directly.

1# PyTorch
2import torch
3x = torch.tensor([1.0, 2.0, 3.0])
4zeros = torch.zeros(3, 4)
5randn = torch.randn(3, 4)

1# TensorFlow
2import tensorflow as tf
3x = tf.constant([1.0, 2.0, 3.0])
4zeros = tf.zeros([3, 4])
5randn = tf.random.normal([3, 4])

Operations such as reshape, concatenation, matrix multiplication, and transpose also have close equivalents. The conceptual model is nearly the same even when naming differs.

Gradient Computation

PyTorch uses autograd through backward, while TensorFlow uses GradientTape.

1# PyTorch
2x = torch.tensor(3.0, requires_grad=True)
3y = x ** 2 + 2 * x + 1
4y.backward()
5print(x.grad)

1# TensorFlow
2x = tf.Variable(3.0)
3with tf.GradientTape() as tape:
4    y = x ** 2 + 2 * x + 1
5grad = tape.gradient(y, x)
6print(grad)

These are the same idea expressed with different control flow.

Model Definition

The closest TensorFlow equivalent to a PyTorch nn.Module is a subclass of tf.keras.Model.

1# PyTorch
2import torch.nn as nn
3
4class MyModel(nn.Module):
5    def __init__(self):
6        super().__init__()
7        self.fc1 = nn.Linear(784, 128)
8        self.fc2 = nn.Linear(128, 10)
9
10    def forward(self, x):
11        x = torch.relu(self.fc1(x))
12        return self.fc2(x)

1# TensorFlow
2import tensorflow as tf
3
4class MyModel(tf.keras.Model):
5    def __init__(self):
6        super().__init__()
7        self.fc1 = tf.keras.layers.Dense(128, activation="relu")
8        self.fc2 = tf.keras.layers.Dense(10)
9
10    def call(self, x):
11        x = self.fc1(x)
12        return self.fc2(x)

TensorFlow also offers Sequential and model.fit, which feel more high-level than the typical PyTorch style, though custom loops are still possible.

Training Loops

A manual TensorFlow loop with GradientTape is the closest equivalent to a PyTorch training loop.

1# PyTorch
2optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
3criterion = nn.CrossEntropyLoss()
4
5for inputs, targets in dataloader:
6    optimizer.zero_grad()
7    outputs = model(inputs)
8    loss = criterion(outputs, targets)
9    loss.backward()
10    optimizer.step()

1# TensorFlow
2optimizer = tf.keras.optimizers.Adam(1e-3)
3loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
4
5for inputs, targets in dataset:
6    with tf.GradientTape() as tape:
7        outputs = model(inputs, training=True)
8        loss = loss_fn(targets, outputs)
9    grads = tape.gradient(loss, model.trainable_variables)
10    optimizer.apply_gradients(zip(grads, model.trainable_variables))

If you want more automation in TensorFlow, model.compile and model.fit sit on top of the same concepts.

Data Loading

PyTorch uses Dataset and DataLoader. TensorFlow uses tf.data.Dataset.

1# PyTorch
2from torch.utils.data import DataLoader, TensorDataset
3
4dataset = TensorDataset(torch.randn(1000, 784), torch.randint(0, 10, (1000,)))
5loader = DataLoader(dataset, batch_size=32, shuffle=True)

1# TensorFlow
2dataset = tf.data.Dataset.from_tensor_slices(
3    (tf.random.normal([1000, 784]), tf.random.uniform([1000], 0, 10, dtype=tf.int32))
4)
5dataset = dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)

The big conceptual difference is that tf.data often encourages a more pipeline-oriented style.

Saving and Loading Models

PyTorch commonly saves state_dict. TensorFlow commonly saves weights or the full Keras model.

# PyTorch
torch.save(model.state_dict(), "model.pth")
model.load_state_dict(torch.load("model.pth"))

1# TensorFlow
2model.save_weights("model.weights.h5")
3model.load_weights("model.weights.h5")
4
5model.save("saved_model_dir")
6model = tf.keras.models.load_model("saved_model_dir")

The conceptual match is easy: save parameters only, or save the full model graph and weights together.

Device Placement and Layout Differences

PyTorch often makes device movement explicit with .to(device). TensorFlow often handles GPU placement automatically.

One practical difference matters a lot in vision code:

• PyTorch commonly uses NCHW
• TensorFlow commonly uses NHWC

This affects convolution input layout and data preprocessing, and it is one of the most frequent migration pitfalls.

Common Pitfalls

A common mistake is assuming names differ but defaults do not. Layout conventions, loss expectations, and serialization styles often differ in meaningful ways.

Another mistake is translating PyTorch custom-loop code directly into TensorFlow while ignoring the higher-level tf.keras tools that may already fit the use case.

Developers also often overlook channel order when porting image models.

Finally, do not treat feature mapping as only an API rename exercise. Framework defaults and ecosystem patterns matter too.

Summary

• TensorFlow and PyTorch have close equivalents for tensors, gradients, layers, training loops, and data loading.
• 'tf.keras.Model plus GradientTape is the closest TensorFlow analogue to PyTorch nn.Module plus autograd.'
• 'tf.data.Dataset is the TensorFlow counterpart to PyTorch dataset loaders.'
• Saving, loading, and device placement all have clear conceptual matches.
• Pay special attention to defaults such as channel ordering when translating code between frameworks.