← All writing

PyTorch

February 8, 202612 min#cs#ai#ml

What even is PyTorch ?

PyTorch is a library for Python programs that facilitates the computations for running deep learning models. It emphasizes flexibility and allows deep learning models to be expressed in idiomatic Python. Much of PyTorch’s ease of use is due to how it implements deferred execution.

Deferred execution means that most exceptions are be raised when the function is called, not when it’s defined. For normal Python, that’s fine, because the interpreter and debuggers have full access to the Python state at the time when the error occurred. Things get tricky when specialized classes that have heavy operator overloading are used, allowing what looks like immediate execution to be deferred under the hood. The way this is achieved is by the idea of computational graph.

image.png

Graph mode is the default mode of operation in TensorFlow 1.0. By contrast, PyTorch sports a define-by-run dynamic graph engine in which the computation graph is built node by node as the code is eagerly evaluated.

image.png

The above figure shows the same calculation running under a dynamic graph engine. The computation is broken into individual expressions, which are greedily evaluated as they’re encountered.

For performance reasons, most of PyTorch is written in C++ and CUDA, a C++-like language from NVIDIA that can be compiled to run with massive parallelism on NVIDIA GPUs. There are ways to run PyTorch directly from C.

The core idea behind PyTorch is multidimensional arrays called tensors.

A tensor is like a NumPy array, but it can run on a GPU and supports automatic differentiation.

import torch

x = torch.tensor([1.0, 2.0, 3.0])
y = x * 2
print(y)

The second core thing that PyTorch provides allows tensors to keep track of the operations performed on them and to compute derivatives of an output with respect to any of its inputs analytically via backpropagation.

Compared with NumPy arrays, PyTorch tensors have a few superpowers, such as the ability to perform fast operations on graphical processing units (GPUs), to distrib ute operations on multiple devices or machines, and to keep track of the graph of computations that created them. All these features are important in implementing a modern deep learning library.

This capability is provided natively by tensors and further refined in torch.autograd. We could argue that by having tensors and the autograd-enabled tensor standard library, PyTorch could be used for more than neural networks, and we’d be correct: PyTorch can be used for physics, rendering, optimization, simulation, modeling, and so on.

image.png

Utilities for data loading and handling can be found in torch.util.data. The two main classes you’ll work with are Dataset, which acts as the bridge between your cus tom data (in whatever format it might be in), and a standardized PyTorch Tensor. The other class you’ll see a lot of is DataLoader, which can spawn child processes to load data from a Dataset in the background so that it’s ready and waiting for the training loop as soon as the loop can use it

In the simplest case, the model will be running the required calculations on the local CPU or on a single GPU, so when the training loop has the data, computation can start immediately. It’s more common, however, to want to use specialized hard ware such as multiple GPUs or to have multiple machines contribute their resources to training the model. In those cases, torch.nn.DataParallel and torch.distrib uted can be employed to leverage the additional hardware available.

When you have results from running your model on the training data, torch.optim provides standard ways of updating the model so that the output starts to more closely resemble the answers specified in the training data. As mentioned earlier, PyTorch defaults to an immediate execution model (eager mode). Whenever an instruction involving PyTorch is executed by the Python inter preter, the corresponding operation is immediately carried out by the underlying C++ or CUDA implementation. As more instructions operate on tensors, more operations are executed by the backend implementation. This process is as fast as it typically can be on the C++ side, but it incurs the cost of calling that implementation through Python. This cost is minute, but it adds up. To bypass the cost of the Python interpreter and offer the opportunity to run mod els independently from a Python runtime, PyTorch also provides a deferred execution model named TorchScript. Using TorchScript, PyTorch can serialize a set of instruc tions that can be invoked independently from Python. You can think of this model as being a virtual machine with a limited instruction set specific to tensor operations. Besides not incurring the costs of calling into Python, this execution mode gives PyTorch the opportunity to Just in Time (JIT) transform sequences of known opera tions into more efficient fused operations. These features are the basis of the produc tion deployment capabilities of PyTorch

Why not a list/array ?

image.png

It’s not unusual for simple Python programs that deal with vectors of numbers, such as the coordinates of a 2D line, to use Python lists to store the vector. This practice can be suboptimal, however, for several reasons: Numbers in Python are full-fledged objects. Whereas a floating-point number might take only 32 bits to be represented on a computer, Python boxes them in a full fledged Python object with reference counting and so on. This situation isn’t a problem if you need to store a small number of numbers, but allocating mil lions of such numbers gets inefficient. Lists in Python are meant for sequential collections of objects. No operations are defined for, say, efficiently taking the dot product of two vectors or summing vectors. Also, Python lists have no way of optimizing the layout of their content in memory, as they’re indexable collections of pointers to Python objects (of any kind, not numbers alone). Finally, Python lists are one-dimensional, and although you can create lists of lists, again, this practice is inefficient. The Python interpreter is slow compared with optimized, compiled code. Performing mathematical operations on large collections of numerical data can be must faster using optimized code written in a compiled, low-level language like C.

How are tensors stored?

Values are allocated in contiguous chunks of memory, managed by torch.Storage instances. A storage is a one-dimensional array of numerical data, such as a contiguous block of memory containing numbers of a given type, perhaps a float or int32. A PyTorch Tensor is a view over such a Storage that’s capable of indexing into that storage by using an offset and per-dimension strides.

image.png

Multiple tensors can index the same storage even if they index into the data differ ently. You can see an example in figure 2.4. In fact, when you requested points[0] in the last snippet, what you got back was another tensor that indexes the same storage as the points tensor, but not all of it and with different dimensionality (1D versus 2D). The underlying memory is allocated only once, however, so creating alternative tensor views on the data can be done quickly, regardless of the size of the data managed by the Storage instance.

image.png

Accessing an element i, j in a 2D tensor results in accessing the storage_offset + stride[0] * i + stride[1] * j element in the storage. The offset will usually be zero; if this tensor is a view into a storage created to hold a larger tensor the offset might be a positive value.

PyTorch Basics — Automatic Differentiation (Part 2)

You said:

“Maybe it's a method that helps optimizing... finding gradients... and derivative matrices faster”

You’re on the right track. Let’s refine it into one clean idea:

Key idea (just one):

Automatic differentiation means PyTorch keeps track of operations on tensors so it can compute gradients for you during backpropagation.

You don’t compute derivatives manually — PyTorch builds a computational graph behind the scenes.

Here’s a tiny example:

import torch

x = torch.tensor(2.0, requires_grad=True)
y = x**2 + 3*x
y.backward()   # compute dy/dx

print(x.grad)  # should be 2*x + 3 = 7

This is the engine that lets neural networks learn.



⭐ STEP 2 — PyTorch Basics: Building a Simple Model

Before we jump into UNet, you need to understand how PyTorch organizes models.

We’ll take this slowly.

Key idea (one at a time):

Every PyTorch model is a class that inherits from nn.Module.

Here’s the smallest possible model:

import torch.nn as nn

class TinyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(10, 5)

    def forward(self, x):
        return self.layer(x)

This structure is the same for UNet — just much bigger.


❓ Your turn

Why do you think PyTorch forces us to define a forward() method inside the model class?

Take a guess — even a simple one.

That’s a good intuition, and now we can turn it into the exact mental model you’ll need when we build UNet.


🌱 PyTorch Basics — Why Every Model Has a forward() Method (Part 7)

You said:

“maybe to let the model train / do a forward pass as more data comes in…”

You’re on the right track. Let’s refine it into the clean, correct idea.

Key idea (one sentence):

forward() defines how your input flows through the layers of your model — it is the blueprint of the forward pass.

PyTorch needs this because:

  • It must know how to compute the output from the input
  • It must record the operations for automatic differentiation
  • It must rebuild the computational graph every time you run a forward pass

So forward() is literally the recipe for your model’s computation.


🧠 Tiny mental picture

Think of your model as a factory:

  • __init__() → builds the machines (layers)
  • forward() → defines how the product moves through the machines

UNet is just a big factory with many downsampling and upsampling steps.


✔️ You now understand:

  • Tensors
  • Computational graph
  • Automatic differentiation
  • What needs gradients
  • How PyTorch models are structured
  • Why forward() exists

This is enough to move to the next essential PyTorch concept before we build UNet.


⭐ STEP 3 — Training Loop Basics (the heart of deep learning)

Before we touch UNet, you need to understand the training loop, because UNet uses the exact same pattern.

Here’s the minimal training loop structure:

for epoch in range(num_epochs):
    for batch in dataloader:
        optimizer.zero_grad()
        outputs = model(batch)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

But don’t worry — we’ll break this down slowly.


❓ Your turn

Let’s start simple:

Why do you think we call optimizer.step() at the end of each batch?

Take a guess — even a short one.

Great — this is exactly the moment where PyTorch starts to “click,” so let’s slow down and build the intuition step by step.

Since you’re in Study Mode, I’ll guide you with one idea at a time and pause for your response.


🌱 PyTorch Training Loop — Understanding optimizer.step() (Part 1)

You said:

“I think it's an optimizer like Adam… but I don't know how this is attached to the model… I don’t exactly know how it works.”

Perfect — this is the right confusion to have.

Let’s clear it up with one simple idea.


Key idea (just one for now):

optimizer.step() updates the model’s weights using the gradients that were just computed.

That’s it.

But to understand how it knows which weights to update, we need to look at how the optimizer is created.


🌱 How the optimizer connects to the model

When you create an optimizer, you pass it the model’s parameters:

optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

This line tells PyTorch:

“Hey optimizer, here are the weights and biases you should update during training.”

So the optimizer stores references to all learnable parameters.

Then the training loop works like this:

  1. Forward pass → compute predictions
  2. Loss → measure how wrong the predictions are
  3. Backward pass → compute gradients for each weight
  4. optimizer.step() → update weights using those gradients
  5. optimizer.zero_grad() → clear old gradients

🧠 Tiny mental picture

Think of it like this:

  • The model has weights
  • Backprop computes gradients for those weights
  • The optimizer is like a worker who says: “Give me the gradients, I’ll adjust the weights.”

optimizer.step() is the moment the worker actually moves the weights.


✔️ You now understand:

  • Why the optimizer exists
  • How it knows which parameters to update
  • What optimizer.step() actually does

This is a huge milestone — now you’re ready to start building real models.


Your turn (quick check)

Why do you think we call optimizer.zero_grad() before computing new gradients?

Take a guess — even a simple one.

You’re thinking about it the right way — now let’s make the idea precise so it becomes automatic for you when we start training UNet.

I’ll give you one idea at a time, then pause for you.


🌱 Training Loop — Why optimizer.zero_grad() Exists (Part 1)

You said:

“But we do change the weights during backpropagation it's not during the forward pass right.

I think optimizer.zero_grad() initializes the gradient with 0s?”

You’re very close. Let’s refine it.

Key idea:

Backpropagation does not change the weights.

It only computes gradients and stores them in param.grad.

The optimizer is the thing that actually changes the weights.

So the sequence is:

  1. forward() → compute predictions
  2. loss.backward() → compute gradients
  3. optimizer.step() → update weights
  4. optimizer.zero_grad() → clear old gradients

🌱 Why do we need optimizer.zero_grad()****? (Part 2)

Here’s the exact reason:

PyTorch accumulates gradients by default.

That means:

  • After the first backward pass, param.grad has some value
  • After the second backward pass, PyTorch adds new gradients on top of the old ones
  • This would completely mess up training

So optimizer.zero_grad() does not “initialize gradients to zero” in a general sense — it clears out the old gradients so the next backward pass starts fresh.


🧠 Tiny mental picture

Think of gradients like water filling a bucket:

  • loss.backward() pours water into the bucket
  • optimizer.step() uses the water to move the weights
  • optimizer.zero_grad() empties the bucket

If you don’t empty the bucket, the next pour overflows and ruins everything.


✔️ You now understand:

  • Backprop computes gradients
  • Optimizer updates weights
  • Gradients accumulate unless cleared
  • optimizer.zero_grad() prevents accumulation

This is a huge conceptual milestone.


Your turn — small check

Why would it be a problem if PyTorch kept adding new gradients on top of old ones during training?

Try to explain it in your own words.