← Back to Autonomy
A Field Guide for the Curious

PyTorch & TensorFlow

The two engines behind modern artificial intelligence — explained in plain words, with practical notes, runnable examples, and a map of where each one shines.

01The Big Picture

PyTorch and TensorFlow are software libraries for building, training, and running machine-learning models — most famously the deep neural networks that power image recognition, language models, recommendation engines, and self-driving systems. They are free and open-source, and together they account for the overwhelming majority of deep-learning work done today.

Here is the key idea: training a neural network requires doing an enormous amount of arithmetic, and — crucially — computing derivatives of that arithmetic so the model can learn from its mistakes. You could write all of that by hand, but it would be painful and slow. These frameworks do three heavy things for you automatically:

They give you a fast container for numbers (a tensor), they run the math on a GPU so it's hundreds of times faster, and they compute the gradients for you through automatic differentiation. You describe the model; the framework handles the calculus and the hardware.

In one sentence Both let you say "here is my model and my data" and they take care of "now do the math fast and figure out how to improve it."

02A Kitchen Analogy

Imagine you're a chef inventing recipes. The tensor is your mixing bowl — it holds the ingredients (numbers) in an organized way: a single value, a row, a grid, or a stack of grids. The GPU is an industrial kitchen with a thousand line cooks who can chop in parallel, instead of one cook doing everything alone. And automatic differentiation is a magical assistant who watches every step you take and can instantly tell you, "if you'd used a pinch less salt, the dish would have scored higher" — which is exactly the feedback a model needs to improve.

PyTorch and TensorFlow are two competing brands of this fully-equipped kitchen. They cook the same dishes; they just lay out the counters differently and hand you slightly different utensils.


03What They Share

Before the differences, it's worth seeing how much overlap there is. Both frameworks provide:

Tensors

An n-dimensional array of numbers — the universal currency of deep learning. A photo is a 3-D tensor (height × width × color); a batch of photos is 4-D. Everything flowing through a model is a tensor.

Autograd (automatic differentiation)

The engine that tracks every operation and computes gradients automatically, so the network knows how to nudge its parameters to reduce error. This is the beating heart of training.

GPU acceleration & optimizers

Both move computation to NVIDIA/AMD GPUs (and increasingly other accelerators) with one line of code, and both ship ready-made optimizers (SGD, Adam, etc.), loss functions, and neural-network layers.

High-level building blocks

Pre-built layers, datasets, model zoos, and friendly wrappers — Keras for TensorFlow, and modules like torch.nn (plus libraries such as PyTorch Lightning) for PyTorch.

04The Training Loop, Drawn

Almost all deep-learning training — in either framework — is the same five-step cycle, repeated thousands of times. Each lap nudges the model a little closer to correct. Here is the whole idea in one picture:

Data a batch of examples Model forward → prediction Loss how wrong? Gradients backward() Optimizer update weights REPEAT EACH STEP ↺
The five-step loop: feed data in, predict, measure how wrong, compute gradients, update — then repeat.

The top row is the forward pass (making a guess); the bottom row is the backward pass (learning from the error). PyTorch makes you write this loop by hand; Keras hides it inside .fit(). Either way, the cycle underneath is identical.

05Meet PyTorch

Created by Meta (Facebook) AI · 2016 · now under the PyTorch Foundation / Linux Foundation

The researcher's favorite

  • Define-by-run. The model is built on the fly as your Python code executes. Loops, conditionals, and printing all work naturally.
  • Feels like normal Python. Easy to read, easy to debug — drop a breakpoint anywhere.
  • Dominant in research. Most new papers and most open-source models (e.g. on Hugging Face) are PyTorch-first.
  • Strong ecosystem. torchvision, torchaudio, Lightning, and a fast-maturing deployment story.

PyTorch's signature trait is that the computation graph is dynamic: it's created freshly each time you run your code. This "eager" style means a model is just ordinary Python that happens to operate on tensors — which is why beginners and researchers tend to find it the gentler on-ramp.

PyTorch — a tensor and a gradient
import torch

# x is something we want to learn; track its gradient
x = torch.tensor([2.0], requires_grad=True)

y = x**2 + 3*x        # define some computation
y.backward()             # autograd computes dy/dx for us

print(x.grad)            # tensor([7.]) -> derivative 2x+3 at x=2
Practical note The requires_grad=True flag is what tells PyTorch to watch a tensor. Call .backward() and every watched tensor gets its gradient filled in automatically. This single idea scales from this two-line example all the way up to billion-parameter language models.

06Meet TensorFlow

Created by Google Brain · released 2015

The production workhorse

  • Battle-tested deployment. Mature tooling to serve models at scale (TensorFlow Serving), on phones and microcontrollers (TF Lite / LiteRT), and in browsers (TensorFlow.js).
  • Keras built in. The clean, high-level Keras API makes standard models very concise.
  • Graph optimization. Can compile your model into a static graph for speed and portability.
  • Broad hardware reach. First-class support for Google's TPUs alongside GPUs and CPUs.

TensorFlow arrived first and originally used a "define-then-run" model: you built a static computation graph, then executed it. That was efficient for deployment but awkward to debug. TensorFlow 2.x changed the default to eager execution — so it now feels much more like PyTorch — while keeping the option to compile to a fast graph with a single decorator.

TensorFlow — the same gradient example
import tensorflow as tf

x = tf.Variable([2.0])         # a learnable variable

with tf.GradientTape() as tape:  # "record" the math
    y = x**2 + 3*x

print(tape.gradient(y, x))    # tf.Tensor([7.]) -> same answer
Practical note Where PyTorch quietly tracks gradients via requires_grad, TensorFlow asks you to record operations inside a GradientTape block. Different spelling, identical idea: watch the math, then ask for the slope.

07A Short History

The "framework wars" shaped how these tools look today. A quick timeline:

2007–2015

Early frameworks — Theano, Torch (Lua), Caffe — prove the idea but are clunky. Google builds an internal system called DistBelief.

2015

TensorFlow 1.0 is open-sourced by Google. Powerful and production-ready, but its static define-then-run graphs are notoriously hard to debug.

2016

PyTorch launches from Meta with dynamic, Pythonic graphs. Researchers love it almost immediately.

2019

TensorFlow 2.0 makes eager execution the default and folds in Keras as the official high-level API — a direct response to PyTorch's ergonomics.

2018–2022

PyTorch overtakes TensorFlow in research publications. Google quietly seeds a third option, JAX, for its own large-scale research.

2022–today

PyTorch moves to the neutral PyTorch Foundation, adds torch.compile for graph-level speed, and grows edge tooling (ExecuTorch). TensorFlow rebrands TF Lite as LiteRT and keeps its grip on production and on-device inference.


08Side by Side

A practical comparison. None of these are absolutes — both frameworks have spent years copying each other's best ideas, so the gaps are smaller than they used to be.

DimensionPyTorchTensorFlow
Maker / yearMeta, 2016Google, 2015
Default styleDynamic (define-by-run)Eager since 2.x, compiles to graph on demand
Feels likePlain PythonPython + Keras conventions
Learning curveGentle, intuitiveGentle via Keras; deeper layers get complex
DebuggingEasy — standard Python toolsEasy in eager mode; trickier once compiled
Research mindshareDominant in papers & new modelsStrong but declining in research
DeploymentRapidly maturing (TorchServe, ExecuTorch, ONNX)Very mature (Serving, LiteRT, TF.js)
Mobile / embeddedExecuTorch, PyTorch MobileLiteRT (formerly TF Lite) — long the standard
Specialty hardwareGPU-centric, growing accelerator supportFirst-class TPU support
High-level APItorch.nn, PyTorch LightningKeras (built in)

09Same Model, Both Ways

A tiny neural network — a single linear layer that learns to map an input to an output — written in each framework. Use the tabs to flip between them and notice how similar the shapes are.

import torch, torch.nn as nn

model = nn.Linear(1, 1)                 # 1 input -> 1 output
loss_fn = nn.MSELoss()
opt = torch.optim.SGD(model.parameters(), lr=0.01)

x = torch.tensor([[1.],[2.],[3.]])
y = torch.tensor([[2.],[4.],[6.]])   # target: y = 2x

for epoch in range(200):
    pred = model(x)
    loss = loss_fn(pred, y)
    opt.zero_grad()                       # reset old gradients
    loss.backward()                       # compute new gradients
    opt.step()                            # update the weights
import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Dense(1, input_shape=(1,))
])
model.compile(optimizer=tf.keras.optimizers.SGD(0.01),
              loss="mse")

x = [[1.],[2.],[3.]]
y = [[2.],[4.],[6.]]            # target: y = 2x

model.fit(x, y, epochs=200, verbose=0)   # training loop is hidden
What this shows PyTorch makes the training loop explicit — you see every step (predict → measure error → get gradients → update). Keras hides the loop behind .fit(), which is wonderfully concise for standard tasks but gives you less to look at when something custom is needed. This trade-off — visible control vs. tidy abstraction — captures the cultural difference between the two camps.

10End-to-End: A Real Model

The toy example above learns a single line. Here's something closer to real life: a small network that learns to recognize handwritten digits (the classic MNIST dataset — 28×28 grayscale images of the numbers 0–9). This is the "hello world" of deep learning, and it shows the full arc: load data, define layers, train.

import torch, torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 1. Data: 28x28 grayscale digits -> tensors
tfm = transforms.ToTensor()
train = datasets.MNIST("data", train=True, download=True, transform=tfm)
loader = DataLoader(train, batch_size=64, shuffle=True)

# 2. Model: flatten -> hidden layer -> 10 class scores
model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(28*28, 128), nn.ReLU(),
    nn.Linear(128, 10),
)
loss_fn = nn.CrossEntropyLoss()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

# 3. Train: one full pass over the data
model.train()
for images, labels in loader:
    preds = model(images)
    loss = loss_fn(preds, labels)
    opt.zero_grad()
    loss.backward()
    opt.step()
import tensorflow as tf

# 1. Data: built-in MNIST, pixels scaled to 0-1
(x, y), _ = tf.keras.datasets.mnist.load_data()
x = x / 255.0

# 2. Model: flatten -> hidden layer -> 10 class scores
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation="relu"),
    tf.keras.layers.Dense(10),
])
model.compile(
    optimizer="adam",
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"],
)

# 3. Train: one full pass over the data
model.fit(x, y, epochs=1, batch_size=64)
Read it as a story Both programs say the same thing in three beats: get the digits, stack a couple of layers, train. The 128-neuron hidden layer with a ReLU is what lets the network learn curved, non-linear patterns — without it, you'd only ever fit a straight line. Run a few more epochs and either model lands around 97% accuracy.

11Saving & Loading

A trained model is worthless if you can't store it and reload it later. This is one of the first real-project needs — and the two frameworks have slightly different philosophies.

# PyTorch favors saving just the learned weights ("state dict")
torch.save(model.state_dict(), "model.pt")

# To load: rebuild the SAME architecture, then pour the weights in
model = build_model()              # your own function/class
model.load_state_dict(torch.load("model.pt"))
model.eval()                       # switch to inference mode
# Keras can save the WHOLE model: architecture + weights + config
model.save("model.keras")

# Load it back in one line, ready to predict
model = tf.keras.models.load_model("model.keras")
predictions = model.predict(x)
Practical note The key difference: PyTorch's default saves only the weights, so you must keep the model's code to rebuild its structure before loading. Keras can bundle everything into one file. PyTorch's approach is more portable across code versions and is the safer habit for long-lived projects; Keras's is more convenient for quick handoffs.

12The Ecosystem & a Word on JAX

Neither framework lives alone — each anchors a constellation of companion libraries that do the real day-to-day work.

Around PyTorch

  • torchvision / torchaudio / torchtext — datasets, transforms, and pretrained models for images, audio, text.
  • Hugging Face Transformers — the de-facto hub for state-of-the-art models, PyTorch-first.
  • PyTorch Lightning — removes boilerplate from the training loop.
  • torch.compile — opt-in graph compilation for speed.

Around TensorFlow

  • Keras — the built-in high-level API.
  • TensorBoard — superb visualization of training (also usable from PyTorch).
  • TFX — end-to-end production ML pipelines.
  • TF Hub / TF Datasets — reusable models and ready-made datasets.

The third framework: JAX

Increasingly you'll hear about JAX, also from Google. It isn't a drop-in replacement — think of it as "NumPy with superpowers": a familiar array API plus composable function transforms (grad for gradients, jit for compilation, vmap for auto-batching), all compiled through XLA to run fast on GPUs and TPUs. It's favored in cutting-edge research and underpins libraries like Flax and Haiku, as well as a good deal of large-scale model training at Google/DeepMind. For most newcomers PyTorch is still the place to start, but JAX is worth knowing exists.

13Who Uses What

Rough, real-world associations (these shift over time, and most big organizations use more than one):

FrameworkCommonly associated withSweet spot
PyTorchMeta, OpenAI, much of academia, the Hugging Face model ecosystem, many autonomy/robotics teamsResearch, rapid prototyping, large language models, anything using the latest published architectures
TensorFlowGoogle products historically, many enterprises with established ML pipelines, on-device and mobile apps via LiteRTProduction serving at scale, mobile/embedded/browser inference, mature MLOps pipelines
JAXGoogle Research / DeepMind and a growing slice of frontier-model labsHigh-performance, large-scale research where compilation and TPUs matter
Takeaway The "best" framework is often just the one your team, your tutorial, or the model you're building on already uses. The concepts transfer almost entirely — learn one well and picking up another takes days, not months.

14Which Should You Use?

A few honest heuristics that hold up well:

Lean PyTorch if…

You're learning deep learning, doing research, prototyping new ideas, or working with the latest published models. Its readability and the size of its research community make it the path of least resistance today.

Lean TensorFlow if…

Your priority is shipping to constrained or massive-scale targets — mobile apps, microcontrollers, browsers, or large serving fleets — or your team and existing codebase already live in the TensorFlow/Keras world.

Quick recommender
Answer three questions for a (lighthearted) suggestion.
1. What's your main goal?
2. Where will the model run?
3. Any existing code or models?

15Getting to Production

Building a model is half the job; running it somewhere real is the other half. This is where the frameworks' personalities matter most.

Cloud & servers

TensorFlow Serving has long been a standard for high-throughput model serving; PyTorch answers with TorchServe and broad support for the open ONNX format, which lets a model trained in one framework run in many runtimes.

Phones, microcontrollers & the edge

This was historically TensorFlow's stronghold via TF Lite (now LiteRT), widely used for on-device inference. PyTorch has closed in with ExecuTorch. A key technique on both sides is quantization — shrinking a model's numbers from 32-bit floats down to 8-bit integers (INT8) so it runs fast and small on limited hardware.

16Common Gotchas

Forgetting to zero gradients (PyTorch)PyTorch accumulates gradients by default. Forget optimizer.zero_grad() in your loop and gradients pile up across steps, quietly wrecking training. It's the single most common beginner bug.
CPU/GPU device mismatchTensors on the CPU can't directly operate with tensors on the GPU. In PyTorch you'll move things with .to(device); forgetting leads to a loud error. Keep your data and model on the same device.
Version 1.x vs 2.x confusion (TensorFlow)Old tutorials use TensorFlow 1.x's session-and-graph style, very different from modern 2.x. If example code mentions tf.Session() or placeholder, it's outdated — find a 2.x version.
Mixing up training and evaluation modeLayers like dropout and batch-norm behave differently during training vs. inference. Remember model.train() / model.eval() in PyTorch (Keras handles this inside fit/predict). Skipping it gives mysteriously bad predictions.

17Mini Glossary

TensorAn n-dimensional array of numbers — the data structure everything flows through.
Autograd / autodiffAutomatic computation of derivatives (gradients) so the model can learn.
GradientThe slope telling each parameter which way (and how much) to change to reduce error.
GPU / TPUSpecialized chips that do the massively parallel math of deep learning fast.
EpochOne full pass over the training dataset.
BatchA small chunk of examples processed together in one step.
Loss functionA number measuring how wrong the model currently is; training minimizes it.
OptimizerThe rule (e.g. SGD, Adam) for updating parameters using the gradients.
ReLU / activationA simple non-linear function between layers that lets networks learn curved patterns.
Eager vs. graphRunning math line-by-line (eager) vs. compiling it into a fixed graph first (graph).
QuantizationShrinking numbers (e.g. 32-bit float → 8-bit int) for smaller, faster models on limited hardware.
ONNXAn open format for sharing trained models between frameworks and runtimes.

18References & Further Reading

  1. PyTorch — Official documentation and tutorials. https://pytorch.org/ · pytorch.org/tutorials
  2. TensorFlow — Official documentation and guides. https://www.tensorflow.org/ · tensorflow.org/tutorials
  3. Keras — High-level API documentation. https://keras.io/
  4. JAX — Documentation. https://jax.readthedocs.io/
  5. Paszke et al., "PyTorch: An Imperative Style, High-Performance Deep Learning Library," NeurIPS 2019. arxiv.org/abs/1912.01703
  6. Abadi et al., "TensorFlow: A System for Large-Scale Machine Learning," OSDI 2016. arxiv.org/abs/1605.08695
  7. MNIST handwritten digit database — Yann LeCun et al. yann.lecun.com/exdb/mnist
  8. LiteRT (formerly TensorFlow Lite) — on-device inference. ai.google.dev/edge/litert
  9. ExecuTorch — PyTorch for edge and mobile. pytorch.org/executorch
  10. ONNX — Open Neural Network Exchange. https://onnx.ai/
  11. Hugging Face — model hub showing framework usage across models. huggingface.co/models

Note: framework features and product names evolve quickly (for example, TF Lite was rebranded LiteRT). When precise version-specific behavior matters, the official docs above are the authoritative source.