Why you need an ML framework

Without a stack you hand-roll matrix multiplies, storing activations for the backward pass, weight updates, batch loading, and checkpointing. With a framework it looks like a library with a clear API: tensor on device, layer, optimizer, dataloader.

Analogy: you can build an engine from bolts and steel, but for a trip you want a car with a gearbox and a dashboard. The framework is the car for the cycle “data → prediction → error → gradient → new weights”.

Inside the stack: tensor, graph, autograd

Typical architecture:

                 ML Framework
                      │
        ┌─────────────┼─────────────┐
        ↓             ↓             ↓
     Tensor        Autograd       Models
        │             │             │
        ↓             ↓             ↓
   CPU / GPU       Gradients      Layers
        │             │             │
        └─────────────┼─────────────┘
                      ↓
                  Optimizer → Training

Tensor — a multidimensional array with dtype, shape, and device. An image is often [C, H, W]; a batch is [B, C, H, W]. Neural nets live on tensors because almost every operation is bulk linear algebra.

Computational graph — a record of which operations produced a tensor (for the backward pass).

Automatic differentiation (autograd) — automatic gradient computation along that graph.

On top: loss function, optimizer (SGD, Adam, …), layers (Linear, Conv2d), a model as a composition of layers, Dataset/DataLoader, device abstraction (cpu / cuda / …), weight serialization.

Training loop:

data → model → prediction → loss → backpropagation → gradients → optimizer → new weights

micrograd and tinygrad: understand, not only call the API

micrograd (Andrej Karpathy’s idea) — a tiny scalar autograd implementation. The goal is clarity, not speed: each node holds a value and a local gradient; backward walks edges in reverse. After that section you see autograd is not magic — it is a careful graph traversal.

tinygrad — a small framework with its own minimalist, transparent-stack philosophy (not “PyTorch lite”). Useful when you want to see how an op reaches the accelerator without miles of industrial abstraction. You rarely pick it first for production CV; for learning how the ML stack is built — yes.

On fingers. A node stores a value and a local contribution to the derivative. If c = a * b, then on the backward pass the contribution to a is proportional to b, and to b proportional to a. The chain a → b → c is successive applications of the chain rule. That is what loss.backward() does in a large stack — on tensors, with memory optimizations.

PyTorch: the practical default

PyTorch became the de facto standard for deep learning in research and much of applied CV/LLM work. Key pieces: torch.Tensor, torch.nn, torch.optim, Dataset/DataLoader, autograd, CUDA, saving (state_dict, .pt).

Minimal loop:

for x, y in dataloader:
    optimizer.zero_grad()
    prediction = model(x)
    loss = criterion(prediction, y)
    loss.backward()
    optimizer.step()

zero_grad clears old gradients; forward computes the prediction; criterion is the error; backward fills .grad; step updates weights. One line loss.backward() on the outside — inside: graph traversal, derivatives, memory, and kernel launches on the chosen device.

For OCR, CRNN, the YOLO ecosystem, and most vision pilots, PyTorch is a sensible start. Weights often move to ONNX for server or browser deployment next.

Saving a model is a separate skill. Usually you save a state_dict (weights), not the whole Python object: easier to version model code and load a checkpoint on another machine. .pt / .pth is a torch ecosystem convention, not a universal “any neural net” format.

TensorFlow, Keras, and JAX

TensorFlow was long the main industrial stack. Today it matters together with Keras, tf.data, model serving, and places where TF infrastructure is already in place. TensorFlow 2.x emphasized eager execution — closer to familiar imperative style.

Keras — high-level model API. Historically it merged with TF as tf.keras; Keras also evolves as a convenient layer over different backends. Simplified:

Keras → high-level model API
TensorFlow → compute and ecosystem (often as backend)

JAX — a different angle: “NumPy + differentiation + JIT + accelerators” in a functional style. jax.numpy, grad, jit, vmap — transformations of functions, not a stateful nn.Module object like classic PyTorch API. Strong in research and on TPU; steeper entry if you come from object-oriented nn.Module.

A simple JAX mental model: write the loss as a pure function of parameters and a batch, then grad(loss_fn) yields the gradient, jit speeds things via compilation, vmap vectorizes over the batch axis. Model state is usually explicit arrays/trees, not hidden buffers inside a layer object. Great for research and unusual transforms; it asks for discipline if you are used to “a layer with .parameters() inside”.

Other names and an important distinction

Brief periphery: MXNet — historically important, less common now; PaddlePaddle — its own ecosystem; OneFlow — niche scenarios; MLX — interesting on Apple Silicon (unified memory, Apple GPU).

ONNX Runtime — primarily an inference engine, not a full training framework in the PyTorch sense:

Training (PyTorch / TF / …) → ONNX → ONNX Runtime → Inference

Separately: training framework ≠ library of ready-made models.

PyTorch          → tool to build and train
Ultralytics YOLO → ready detection pipeline on top of the ecosystem
HF Transformers  → models and pipelines on top of frameworks
timm / Diffusers → model zoos and utilities

Confusing “we’ll install YOLO” with “we picked a framework” is a common beginner mistake.

What to pick for the task

This is not a “who is best” ranking — it is matching the tool to the context.

CPU, GPU, and TPU: devices and the hardware contract

Stack placement:

Framework → CPU | GPU (CUDA / Apple / …) | TPU

• CPU — always available; often enough for learning models and debugging.
• NVIDIA CUDA — the main accelerator for PyTorch/TensorFlow in production and research.
• Apple GPU — Metal / MPS in PyTorch, plus interest in MLX on unified memory.
• AMD — depends on the stack (ROCm and builds); check support for your framework version.
• TPU — a strength of JAX/TensorFlow in Google cloud scenarios.

Accelerator availability depends not only on the framework name but on drivers, builds, and backend. device = "cuda" is a contract with your environment, not a spell. Moving tensors and the model to a device (to(device)) must stay consistent — otherwise you get an error or a quiet speed loss.

Training vs inference and the path to hardware

Training: forward, loss, backward, gradients, optimizer.
Inference: load weights, feed data, get output — no training.

TRAINING:  data → model → loss → backprop → weights
INFERENCE: data → model → result

From Python to GPU:

Python → Framework API → tensor ops → backend → CUDA/Metal/TPU runtime → hardware

A line like x @ weights can unfold into a huge amount of work on the GPU — so a “simple API” does not mean a “simple implementation”.

Browser link:

PyTorch (training) → ONNX → Transformers.js / ORT Web + WebGPU (inference in the tab)

Client-side details — WebGPU + Transformers.js; export — ONNX: from PyTorch to Runtime.

Comparison and a beginner roadmap

Learning route:

Python → NumPy → network basics → micrograd → PyTorch → CNN → CRNN → Transformers → LLM

Vision branch: PyTorch → YOLO / detection → OCR.

A practical mini-project

Take handwritten digits (MNIST or your own set) and walk the loop once by hand:

1. Load data and build a DataLoader with batches.
2. Define a small CNN (Conv → activation → pooling → linear layer).
3. Write the training loop: zero_grad → forward → loss → backward → step.
4. Measure accuracy on a held-out set.
5. Save a state_dict to .pt.
6. In a separate script, load weights and run inference with no training.
7. Mentally map the same steps to micrograd/tinygrad: where the graph is, where the gradient is, where API “magic” hides.

Hands-on digits on the site: CNN → TrOCR.

Common mistakes

• Treating PyTorch as “the neural net”, not the framework.
• Confusing the framework with a ready model zoo (YOLO/Transformers).
• Not understanding tensor shape and batch dimension.
• Copying backward() without a mental picture of the graph.
• Assuming GPU is mandatory for every learning exercise.
• Mixing up training and inference; expecting .pt to be universal for any runtime.
• Picking device at random without checking drivers.

FAQ

Further reading

• AI hardware
• PyTorch vs TensorFlow
• ONNX: from PyTorch to Runtime
• WebGPU + Transformers.js
• YOLO detection engineering
• Compact CRNN
• Handwritten digits
• Three layers of the LLM stack

Conclusion

A neural net is a mathematical model. A framework is infrastructure around tensors, gradients, and devices. micrograd and tinygrad teach how it works; PyTorch covers most applied training; TensorFlow/Keras and JAX remain important branches for ecosystem and accelerators; inference often lives separately — in ONNX Runtime or even in the browser. Picking a tool answers “what task, what hardware, who maintains it”, not a hunt for a fashionable logo.

Concepts (tensor, grad) → Frameworks → Models (YOLO, CRNN, LLM)
                              ↓
                           Training → GPU/TPU/CPU
                              ↓
                           Inference → Server | Browser (WebGPU)