TensorFlow vs PyTorch — Comparison Guide for Beginners

# first_pytorch.py — a small model, end to end
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 1. Data
x = torch.randn(1000, 10)
y = (x.sum(dim=1) > 0).long()
ds = TensorDataset(x, y)
dl = DataLoader(ds, batch_size=32, shuffle=True)

# 2. Model
model = nn.Sequential(
    nn.Linear(10, 32),
    nn.ReLU(),
    nn.Linear(32, 2),
)

# 3. Loss + optimiser
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-3)

# 4. Training loop — explicit, readable
for epoch in range(5):
    for xb, yb in dl:
        optimizer.zero_grad()
        logits = model(xb)
        loss = loss_fn(logits, yb)
        loss.backward()
        optimizer.step()
    print(f"epoch {epoch}: loss={loss.item():.4f}")

# 5. Save
torch.save(model.state_dict(), "model.pt")

# Why beginners like this:
# - The training loop is explicit. Nothing is hidden.
# - You can print any tensor at any step.
# - If it breaks, you get a real Python traceback.
#
# The cost:
# - You write the loop yourself (verbose for complex work)
# - You have to remember to zero_grad, backward, step
# - Production tooling (serving, monitoring) is
#   thinner than TF's

# debug_pytorch.py — why PyTorch is easier to debug

import torch

x = torch.randn(4, 3)
print(x)                       # ✓ works, prints a tensor
print(x.shape)                 # ✓ works
print(x[0])                    # ✓ works
print(x.mean(dim=1))           # ✓ works

# You can inspect the gradient like any other value
x = torch.randn(3, requires_grad=True)
y = (x ** 2).sum()
y.backward()
print(x.grad)                  # ✓ works

# And you can drop a breakpoint anywhere
# import pdb; pdb.set_trace()

# What PyTorch gives you:
# ✓ Every tensor is a first-class value
# ✓ Every operation runs immediately
# ✓ Breakpoints, prints, and IDE inspection
#   just work

# What TensorFlow 1.x made hard:
# ✗ You built a graph, then ran it in a session
# ✗ print() didn't show values, it showed
#   symbolic tensors
# ✗ Errors surfaced at session.run, not
#   where the bug was

# TensorFlow 2.x fixed this with eager mode.
# But PyTorch started there, and the ecosystem
# around it (tutorials, papers, Hugging Face)
# stayed PyTorch-first.
#
# For a beginner in 2026, the debugging
# experience is a genuine reason to start
# with PyTorch.

# first_tensorflow.py — the same model with Keras
import tensorflow as tf
from tensorflow.keras import layers, models

# 1. Data
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 28*28) / 255.0

# 2. Model — Keras is genuinely nice
model = models.Sequential([
    layers.Input(shape=(28*28,)),
    layers.Dense(128, activation="relu"),
    layers.Dense(10, activation="softmax"),
])

model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)

# 3. Train — one line
model.fit(x_train, y_train, epochs=3, batch_size=64)

# 4. Save (Keras format in 2026)
model.save("model.keras")

# Why beginners might like this:
# ✓ Less code than PyTorch for standard models
# ✓ Keras hides the training loop
# ✓ Fine for a first project
#
# The trade-off:
# ✗ Less control over the training loop
# ✗ Debugging a custom loop is harder
# ✗ The ecosystem (papers, tutorials) is
#   now PyTorch-first for new work

# Keras is still a fine choice for:
# - Standard image classifiers
# - Quick experiments
# - Teaching an intro course
# - Teams already on TF

# export_tflite.py — TF's real edge story

import tensorflow as tf

# Load a trained model
model = tf.keras.models.load_model("model.keras")

# Convert to TFLite — this runs on mobile
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]  # quantise
tflite_model = converter.convert()

with open("model.tflite", "wb") as f:
    f.write(tflite_model)

print(f"TFLite size: {len(tflite_model) / 1024:.1f} KB")

# Why TFLite is genuinely TensorFlow's advantage:
# ✓ Deployed on billions of Android devices
# ✓ Small model size (quantisation built-in)
# ✓ Runs offline, on-device
# ✓ Standard for mobile ML teams

# What PyTorch does instead:
# - PyTorch Mobile exists but has less adoption
# - Most teams use ONNX → TFLite for mobile
# - Or export to CoreML for iOS
#
# ONNX is the neutral ground:
# - Train in PyTorch
# - Export to ONNX
# - Convert ONNX → TFLite or CoreML
# - This works but has more steps
#
# If your primary target is a mobile app
# shipping on Android, TFLite is the
# shortest path — and that's the one place
# TensorFlow is still the best choice.

# PyTorch — explicit loop
import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Sequential(
    nn.Linear(10, 32),
    nn.ReLU(),
    nn.Linear(32, 2),
)

loss_fn = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-3)

for epoch in range(5):
    for xb, yb in dl:
        optimizer.zero_grad()
        logits = model(xb)
        loss = loss_fn(logits, yb)
        loss.backward()
        optimizer.step()

# Line count: 8 (training loop)
# Reads like: "for each batch, do this"
# Debugging: print(loss) anywhere, works
# Customisation: anything you want

# TensorFlow — Keras handles the loop
import tensorflow as tf
from tensorflow.keras import layers, models

model = models.Sequential([
    layers.Input(shape=(10,)),
    layers.Dense(32, activation="relu"),
    layers.Dense(2, activation="softmax"),
])

model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)

model.fit(x_train, y_train, epochs=5, batch_size=32)

# Line count: 3 (fit call)
# Reads like: "fit the model to this data"
# Debugging: use tf.debugging or
#            model.evaluate, less natural
# Customisation: possible, but you need
#                to know the callbacks API

Same concept, different vocabulary:

PYTORCH                 TENSORFLOW (Keras)
---------------------   ---------------------
nn.Module               layers.Layer
model(x)                model(x)
loss.backward()         gradienttape (auto in fit)
optimizer.step()        optimizer.apply_gradients
optimizer.zero_grad()   (handled by fit)
for epoch in ...        model.fit(...)
loss_fn(logits, y)      model.compile(loss=...)
DataLoader              tf.data.Dataset
torch.save              model.save
state_dict              weights

Which style is better for a beginner?
  PyTorch: teaches how training works because
           everything is explicit
  Keras:   gets you to a working model faster

The honest answer:
  If you're learning deep learning, PyTorch
  is better. You write the loop once and
  understand it forever.
  If you're shipping a standard model
  quickly and don't care about the internals,
  Keras is fine.

What interviews actually test:
  - "What does zero_grad() do?"
  - "What is the difference between
    model.eval() and torch.no_grad()?"
  - "How does backpropagation work under
    the hood?"
  All three are easier to answer if you've
  written the loop yourself in PyTorch.

# export_onnx.py — the neutral deployment path
import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(10, 32),
    nn.ReLU(),
    nn.Linear(32, 2),
)
model.load_state_dict(torch.load("model.pt"))
model.eval()

# Export to ONNX
dummy_input = torch.randn(1, 10)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["features"],
    output_names=["logits"],
    dynamic_axes={"features": {0: "batch"}},
    opset_version=17,
)

# ONNX runs everywhere:
# ✓ onnxruntime on servers (fast, no PyTorch needed)
# ✓ onnxruntime-web in browsers
# ✓ convert → TFLite for mobile
# ✓ convert → CoreML for iOS
# ✓ converts → TensorRT for NVIDIA GPUs

# Why this matters:
# - Trained in PyTorch (good for research)
# - Deployed in ONNX (good for portability)
# - Converted to mobile (good for edge)
# One pipeline, best tool at each step.

# Reality check:
# - Most ONNX exports of standard models
#   work cleanly
# - Some custom ops require a plugin
# - Always test the exported model against
#   the original on a real batch before shipping

# export_tflite.py — TF's mobile path
import tensorflow as tf

model = tf.keras.models.load_model("model.keras")

converter = tf.lite.TFLiteConverter.from_keras_model(model)

# Optional: quantise to int8 — 4x smaller
converter.optimizations = [tf.lite.Optimize.DEFAULT]

tflite_model = converter.convert()
with open("model.tflite", "wb") as f:
    f.write(tflite_model)

# What this gives you:
# - A single .tflite file
# - Runs on Android, iOS, embedded
# - No Python runtime needed
# - Used by billions of devices

# What you lose:
# - Some custom ops aren't supported
# - Debugging is harder after conversion
# - Quantisation can drop accuracy
#   (measure this)

# The deployment decision tree:
# 1. Server only? → FastAPI + PyTorch or ONNX
# 2. Mobile app? → TFLite (TF) or ONNX → TFLite
# 3. Browser? → TF.js or ONNX.js
# 4. Edge device (Raspberry Pi, IoT)? →
#    TFLite or ONNX Runtime

# The honest summary:
# - PyTorch wins the training experience
# - TF wins the edge story (today)
# - ONNX lets you use both