Skip to content

16. PyTorch Fundamentals and Tensors

Learn the foundations of PyTorch, understand tensors, tensor operations, device management, automatic differentiation, neural network modules, parameters, GPU acceleration, and the core building blocks required to develop Deep Learning models using PyTorch.


๐ŸŽฏ Learning Objectives

After completing this chapter, you will be able to:

  • Explain what PyTorch is
  • Understand the architecture of PyTorch
  • Understand PyTorch tensors
  • Understand tensor dimensions, shape, rank, and data types
  • Create tensors using different methods
  • Perform tensor indexing and slicing
  • Perform tensor arithmetic and matrix operations
  • Understand broadcasting
  • Reshape, transpose, squeeze, and unsqueeze tensors
  • Understand tensor devices
  • Move tensors between CPU and GPU
  • Understand tensor memory and performance considerations
  • Understand requires_grad
  • Understand the role of tensors in automatic differentiation
  • Understand torch.nn
  • Build neural network components using nn.Module
  • Understand parameters and trainable variables
  • Build a basic PyTorch neural network
  • Understand the PyTorch training workflow
  • Understand the difference between training and inference
  • Save and load PyTorch models
  • Understand PyTorch GPU acceleration
  • Understand the relationship between tensors, models, autograd, and optimizers
  • Prepare for PyTorch datasets, DataLoaders, and custom training loops

๐Ÿ“– Overview

PyTorch is an open-source Deep Learning framework widely used for:

  • Neural network development
  • Computer Vision
  • Natural Language Processing
  • Generative AI
  • Reinforcement Learning
  • Research
  • Production Machine Learning

The core PyTorch ecosystem can be viewed as:

PyTorch
   โ”‚
   โ”œโ”€โ”€ Tensors
   โ”‚
   โ”œโ”€โ”€ Autograd
   โ”‚
   โ”œโ”€โ”€ torch.nn
   โ”‚
   โ”œโ”€โ”€ Optimizers
   โ”‚
   โ”œโ”€โ”€ Data Utilities
   โ”‚
   โ”œโ”€โ”€ GPU / CUDA
   โ”‚
   โ””โ”€โ”€ Model Serialization

๐Ÿง  What Is PyTorch?

PyTorch is a tensor-based Deep Learning framework that provides:

Tensor Computation
+
Automatic Differentiation
+
Neural Network APIs
+
Optimization
+
GPU Acceleration
+
Data Loading

A typical PyTorch Deep Learning workflow is:

flowchart LR

    DATA["Data"]

    TENSOR["PyTorch Tensors"]

    MODEL["Neural Network"]

    LOSS["Loss"]

    AUTOGRAD["Autograd"]

    OPT["Optimizer"]

    UPDATE["Updated Parameters"]

    DATA --> TENSOR
    TENSOR --> MODEL
    MODEL --> LOSS
    LOSS --> AUTOGRAD
    AUTOGRAD --> OPT
    OPT --> UPDATE
    UPDATE --> MODEL

๐Ÿง  PyTorch vs TensorFlow

Both frameworks provide the fundamental capabilities required for Deep Learning.

Capability PyTorch TensorFlow
Tensor Operations torch.Tensor tf.Tensor
Neural Networks torch.nn tf.keras
Automatic Differentiation torch.autograd tf.GradientTape
Optimizers torch.optim tf.keras.optimizers
Data Pipeline Dataset / DataLoader tf.data
GPU CUDA / device APIs CUDA / device APIs
Model Definition nn.Module Keras Model
Training Custom / higher-level tooling model.fit() / custom
Ecosystem PyTorch ecosystem TensorFlow ecosystem

The important point is that both frameworks implement similar Deep Learning concepts using different APIs and abstractions.


๐Ÿง  The PyTorch Mental Model

A useful mental model is:

Tensor
  โ†“
Model
  โ†“
Prediction
  โ†“
Loss
  โ†“
Autograd
  โ†“
Gradients
  โ†“
Optimizer
  โ†“
Parameter Update

This loop is repeated across batches and epochs.


๐Ÿ”ข What Is a PyTorch Tensor?

A PyTorch tensor is a multidimensional data structure used for numerical computation.

Examples:

Scalar
Vector
Matrix
3D Tensor
4D Tensor
5D Tensor

Deep Learning models operate primarily on tensors.


๐Ÿ“ Tensor Rank

Tensor rank represents the number of dimensions.

Rank 0 โ€” Scalar

x = torch.tensor(5)

Conceptually:

5

Rank 1 โ€” Vector

x = torch.tensor(
    [1, 2, 3]
)

Conceptually:

[1, 2, 3]

Rank 2 โ€” Matrix

x = torch.tensor(
    [
        [1, 2],
        [3, 4]
    ]
)

Conceptually:

[
    [1, 2],
    [3, 4]
]

Rank 3 Tensor

x = torch.tensor(
    [
        [
            [1, 2],
            [3, 4]
        ],
        [
            [5, 6],
            [7, 8]
        ]
    ]
)

๐Ÿง  Tensor Dimensions

A tensor can be represented as:

Shape
 โ†“
(D1, D2, D3, ...)

For example:

(32, 224, 224, 3)

may represent:

Batch Size = 32
Height     = 224
Width      = 224
Channels   = 3
flowchart LR

    B["Batch<br>32"]

    H["Height<br>224"]

    W["Width<br>224"]

    C["Channels<br>3"]

    B --> H
    H --> W
    W --> C

๐Ÿ Importing PyTorch

import torch

For neural networks:

import torch.nn as nn

For optimization:

import torch.optim as optim

๐Ÿงช Creating Tensors

From Python Lists

x = torch.tensor(
    [1, 2, 3]
)

Matrix

x = torch.tensor(
    [
        [1, 2],
        [3, 4]
    ]
)

Random Tensor

x = torch.rand(
    3,
    4
)

Normal Distribution

x = torch.randn(
    3,
    4
)

Zeros

x = torch.zeros(
    3,
    4
)

Ones

x = torch.ones(
    3,
    4
)

Empty Tensor

x = torch.empty(
    3,
    4
)

empty() allocates memory without initializing the values to a meaningful default.


๐Ÿง  Tensor Creation Summary

Function Purpose
torch.tensor() Create from existing data
torch.zeros() Create zeros
torch.ones() Create ones
torch.rand() Uniform random values
torch.randn() Normal random values
torch.empty() Uninitialized tensor
torch.arange() Sequence of values
torch.linspace() Evenly spaced values

๐Ÿ”ข torch.arange()

x = torch.arange(
    0,
    10
)

Result:

[0, 1, 2, ..., 9]

๐Ÿ“ torch.linspace()

x = torch.linspace(
    0,
    1,
    5
)

Conceptually:

[0.00, 0.25, 0.50, 0.75, 1.00]

๐Ÿ” Inspecting a Tensor

x = torch.randn(
    32,
    784
)

print(x)
print(x.shape)
print(x.ndim)
print(x.dtype)
print(x.device)

Typical information:

Shape
Number of Dimensions
Data Type
Device

๐Ÿง  Tensor Shape

For:

x = torch.randn(
    32,
    784
)

the shape is:

(32, 784)

This commonly means:

32 samples
784 features

๐Ÿง  Tensor ndim

x.ndim

returns the number of dimensions.

Example:

x = torch.randn(
    32,
    224,
    224,
    3
)

print(x.ndim)

Result:

4

๐Ÿง  Tensor Data Types

Common PyTorch data types include:

torch.float32
torch.float64
torch.float16
torch.bfloat16
torch.int32
torch.int64
torch.uint8
torch.bool

For many Deep Learning workloads:

float32

is a common default.


๐Ÿ”„ Converting Data Types

x = torch.tensor(
    [1, 2, 3]
)

x = x.float()

Or:

x = x.to(
    torch.float32
)

๐Ÿง  Why Data Type Matters

Data type affects:

Memory Usage
Computation Speed
Numerical Precision
GPU Performance

For example:

float32
   โ†“
Higher Precision

float16 / bfloat16
   โ†“
Lower Memory
Potentially Faster Training

Mixed precision is explored further in advanced training and optimization topics.


๐Ÿ”ข Tensor Indexing

x = torch.tensor(
    [
        [10, 20, 30],
        [40, 50, 60]
    ]
)

Access the first row:

x[0]

Access the first element:

x[0, 0]

Access the second row, third element:

x[1, 2]

โœ‚๏ธ Tensor Slicing

x[:, 0]

selects the first column.

x[0, :]

selects the first row.

x[:, 1:3]

selects columns 1 and 2.


๐Ÿงฎ Tensor Arithmetic

a = torch.tensor(
    [1, 2, 3]
)

b = torch.tensor(
    [4, 5, 6]
)

print(a + b)
print(a - b)
print(a * b)
print(a / b)

These operations are element-wise.


โœ–๏ธ Matrix Multiplication

Matrix multiplication is fundamental to neural networks.

a = torch.tensor(
    [
        [1.0, 2.0],
        [3.0, 4.0]
    ]
)

b = torch.tensor(
    [
        [5.0, 6.0],
        [7.0, 8.0]
    ]
)

result = torch.matmul(
    a,
    b
)

You can also use:

result = a @ b

Mathematically:

[ C=AB ]


๐Ÿงฎ Dot Product

For vectors:

a = torch.tensor(
    [1.0, 2.0, 3.0]
)

b = torch.tensor(
    [4.0, 5.0, 6.0]
)

result = torch.dot(
    a,
    b
)

๐Ÿ”„ Reshaping

PyTorch provides:

reshape()

Example:

x = torch.arange(
    12
)

x = x.reshape(
    3,
    4
)

Result:

3 ร— 4

๐Ÿง  view()

view() can reshape tensors when the underlying memory layout permits it.

x = torch.arange(
    12
)

y = x.view(
    3,
    4
)

In modern PyTorch code, reshape() is often more convenient because it can handle non-contiguous tensors by creating a copy when necessary.


๐Ÿ”„ Flattening

x = torch.randn(
    32,
    28,
    28
)

flat = x.reshape(
    32,
    -1
)

The result is:

32 ร— 784

This is commonly used when transitioning from image feature maps to fully connected layers.


โ†”๏ธ Transpose

x = torch.randn(
    3,
    4
)

y = x.T

For more general dimensions:

y = x.transpose(
    0,
    1
)

๐Ÿงฉ permute()

permute() changes the ordering of dimensions.

Example:

x = torch.randn(
    32,
    224,
    224,
    3
)

y = x.permute(
    0,
    3,
    1,
    2
)

Shape changes from:

Batch ร— Height ร— Width ร— Channels

to:

Batch ร— Channels ร— Height ร— Width

This is particularly important because many PyTorch vision layers commonly use channel-first tensor layouts.


๐Ÿง  Tensor Layout

A common PyTorch image tensor format is:

N ร— C ร— H ร— W

where:

N = Batch
C = Channels
H = Height
W = Width

For example:

32 ร— 3 ร— 224 ร— 224

๐Ÿ–ผ๏ธ Image Tensor Pipeline

flowchart LR

    IMAGE["Image"]

    LOAD["Load"]

    TENSOR["Tensor"]

    FORMAT["N ร— C ร— H ร— W"]

    CNN["CNN"]

    IMAGE --> LOAD
    LOAD --> TENSOR
    TENSOR --> FORMAT
    FORMAT --> CNN

๐Ÿ”„ unsqueeze()

unsqueeze() adds a dimension.

x = torch.tensor(
    [1, 2, 3]
)

y = x.unsqueeze(
    0
)

Shape:

(3)

becomes:

(1, 3)

๐Ÿ”„ squeeze()

squeeze() removes dimensions of size 1.

x = torch.randn(
    1,
    3,
    1
)

y = x.squeeze()

๐Ÿง  Broadcasting

PyTorch supports broadcasting for compatible shapes.

x = torch.tensor(
    [
        [1.0, 2.0],
        [3.0, 4.0]
    ]
)

y = torch.tensor(
    [10.0, 20.0]
)

result = x + y

Conceptually:

[1, 2]    [10, 20]
[3, 4] +  [10, 20]

Result:

[11, 22]
[13, 24]

๐Ÿง  Tensor Reduction

Common reduction operations include:

x.mean()
x.sum()
x.max()
x.min()

Example:

x = torch.tensor(
    [1.0, 2.0, 3.0]
)

print(
    x.mean()
)

๐Ÿ“Š Reduction Along a Dimension

x = torch.tensor(
    [
        [1.0, 2.0],
        [3.0, 4.0]
    ]
)

row_mean = x.mean(
    dim=1
)

column_mean = x.mean(
    dim=0
)

Understanding dimensions is critical when building neural networks.


๐Ÿง  Device Management

PyTorch tensors can live on different devices.

Common examples:

CPU
CUDA GPU
MPS

A tensor's device can be inspected using:

x.device

๐Ÿ–ฅ๏ธ CPU Tensor

x = torch.tensor(
    [1.0, 2.0, 3.0]
)

print(
    x.device
)

Typical result:

cpu

๐Ÿš€ GPU Availability

For NVIDIA CUDA:

torch.cuda.is_available()

Example:

if torch.cuda.is_available():

    print(
        "CUDA GPU available"
    )

๐Ÿง  Selecting a Device

A common pattern is:

device = torch.device(
    "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

Then:

x = torch.randn(
    32,
    784
).to(device)

๐Ÿง  Device-Agnostic Code

A production-friendly approach is:

device = torch.device(
    "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

model = model.to(
    device
)

x = x.to(
    device
)

This allows the same code to run on:

GPU
or
CPU

without hardcoding one environment.


๐Ÿง  CPU โ†’ GPU

x = torch.randn(
    1000,
    1000
)

if torch.cuda.is_available():

    x = x.to(
        "cuda"
    )

๐Ÿง  GPU โ†’ CPU

x_cpu = x.to(
    "cpu"
)

When converting tensors to NumPy:

array = x_cpu.numpy()

For a GPU tensor:

array = x.detach().cpu().numpy()

โš  Device Mismatch

Model and input tensors generally need to be on compatible devices.

Incorrect:

Model โ†’ GPU
Input โ†’ CPU

This can produce runtime errors.

Correct:

Model โ†’ GPU
Input โ†’ GPU
flowchart LR

    MODEL["Model"]

    INPUT["Input"]

    DEVICE["Same Device"]

    MODEL --> DEVICE
    INPUT --> DEVICE

๐Ÿง  GPU Memory

GPU memory is consumed by:

Model Parameters
+
Gradients
+
Optimizer State
+
Activations
+
Input Batches

Therefore, GPU memory usage can increase significantly with:

Larger Models
Larger Batch Sizes
Longer Sequences
Higher Resolution Images
Larger Activations

๐Ÿง  PyTorch Autograd

PyTorch provides automatic differentiation through:

torch.autograd

The key concept is:

requires_grad=True

Example:

x = torch.tensor(
    3.0,
    requires_grad=True
)

๐Ÿงฎ Automatic Differentiation

Suppose:

[ y=x^2 ]

Then:

[ \frac{dy}{dx}=2x ]

PyTorch can calculate this automatically.


๐Ÿงช Basic Autograd Example

import torch


x = torch.tensor(
    3.0,
    requires_grad=True
)

y = x ** 2

y.backward()

print(
    x.grad
)

Result:

6

๐Ÿง  Autograd Workflow

flowchart TD

    X["Input Tensor<br>requires_grad=True"]

    FORWARD["Forward Computation"]

    LOSS["Output / Loss"]

    BACK["backward()"]

    GRAD["Gradients"]

    X --> FORWARD
    FORWARD --> LOSS
    LOSS --> BACK
    BACK --> GRAD
    GRAD --> X

๐Ÿง  Computational Graph

PyTorch tracks operations involving tensors that require gradients.

For:

y = x * x

the computational graph conceptually becomes:

x
โ”‚
โ”œโ”€โ”€ ร—
โ”‚
โ””โ”€โ”€ x
    โ†“
    y

During:

y.backward()

PyTorch traverses the graph backward to compute gradients.


๐Ÿง  Gradient Accumulation

PyTorch gradients accumulate by default.

Example:

x = torch.tensor(
    2.0,
    requires_grad=True
)

y = x ** 2

y.backward()

print(
    x.grad
)

If another backward pass is performed without clearing the gradient, the gradient can accumulate.

Therefore, training loops commonly reset gradients.


๐Ÿงน Clearing Gradients

With an optimizer:

optimizer.zero_grad()

Then:

loss.backward()

Then:

optimizer.step()

The standard sequence is:

zero_grad()
    โ†“
forward
    โ†“
loss
    โ†“
backward()
    โ†“
step()

๐Ÿง  PyTorch Neural Networks

PyTorch provides:

torch.nn

for neural network components.

Common modules include:

nn.Linear
nn.Conv2d
nn.ReLU
nn.Dropout
nn.BatchNorm2d
nn.MaxPool2d
nn.LSTM
nn.GRU
nn.Embedding

๐Ÿงฑ nn.Module

The fundamental abstraction for neural network models is:

nn.Module

A model generally:

class MyModel(
    nn.Module
):

๐Ÿงช Basic PyTorch Model

import torch
import torch.nn as nn


class SimpleNetwork(
    nn.Module
):

    def __init__(self):

        super().__init__()

        self.fc1 = nn.Linear(
            784,
            128
        )

        self.fc2 = nn.Linear(
            128,
            64
        )

        self.output = nn.Linear(
            64,
            10
        )

    def forward(
        self,
        x
    ):

        x = torch.relu(
            self.fc1(x)
        )

        x = torch.relu(
            self.fc2(x)
        )

        return self.output(x)

๐Ÿง  PyTorch Model Architecture

flowchart LR

    INPUT["784 Features"]

    FC1["Linear 128"]

    RELU1["ReLU"]

    FC2["Linear 64"]

    RELU2["ReLU"]

    OUT["Linear 10"]

    INPUT --> FC1
    FC1 --> RELU1
    RELU1 --> FC2
    FC2 --> RELU2
    RELU2 --> OUT

๐Ÿง  Why nn.Module?

nn.Module automatically manages:

  • Parameters
  • Child modules
  • Model hierarchy
  • Device movement
  • Training/inference modes
  • State dictionaries

For example:

model.parameters()

returns trainable parameters.


๐Ÿ” Inspecting Model Parameters

model = SimpleNetwork()

for name, parameter in model.named_parameters():

    print(
        name,
        parameter.shape
    )

Typical parameters include:

fc1.weight
fc1.bias
fc2.weight
fc2.bias
output.weight
output.bias

๐Ÿงฎ Parameter Count

For:

nn.Linear(784, 128)

the parameter count is:

[ 784\times128+128 ]

This includes:

Weights
+
Bias

๐Ÿง  forward()

The forward() method defines the forward computation.

Example:

def forward(
    self,
    x
):

    x = self.fc1(
        x
    )

    x = torch.relu(
        x
    )

    return self.output(
        x
    )

You normally invoke the model as:

output = model(
    x
)

rather than calling:

model.forward(x)

directly.


๐Ÿง  Model Call Flow

flowchart TD

    INPUT["Input"]

    MODEL["model(x)"]

    CALL["nn.Module Call"]

    FORWARD["forward(x)"]

    OUTPUT["Output"]

    INPUT --> MODEL
    MODEL --> CALL
    CALL --> FORWARD
    FORWARD --> OUTPUT

The nn.Module call mechanism also supports hooks and other framework behavior.


๐Ÿง  nn.Linear

The PyTorch equivalent of a fully connected layer is:

nn.Linear(
    in_features,
    out_features
)

Mathematically:

[ y=xW^T+b ]

Example:

layer = nn.Linear(
    784,
    128
)

๐Ÿง  Activation Functions

PyTorch provides activation functions such as:

nn.ReLU()
nn.Sigmoid()
nn.Tanh()
nn.GELU()
nn.Softmax()

Example:

self.relu = nn.ReLU()

or:

x = torch.relu(
    x
)

๐Ÿงช Using nn.Sequential

PyTorch also provides a convenient sequential model API.

model = nn.Sequential(

    nn.Linear(
        784,
        128
    ),

    nn.ReLU(),

    nn.Linear(
        128,
        64
    ),

    nn.ReLU(),

    nn.Linear(
        64,
        10
    )
)

This is conceptually similar to Keras Sequential.


๐Ÿง  Sequential vs Custom nn.Module

Approach Best Use
nn.Sequential Simple linear stacks
Custom nn.Module Complex architectures
Functional Tensor Operations Specialized computations

For architectures involving:

Branches
Skip Connections
Multiple Inputs
Multiple Outputs
Custom Logic

a custom nn.Module is generally more appropriate.


๐Ÿง  Training Mode and Evaluation Mode

PyTorch models have two important modes:

model.train()

and:

model.eval()

Training mode enables training-specific behavior such as:

Dropout
Batch Normalization behavior

Evaluation mode switches the model to inference behavior.


๐Ÿงช Training Mode

model.train()

predictions = model(
    x
)

๐Ÿงช Evaluation Mode

model.eval()

with torch.no_grad():

    predictions = model(
        x
    )

๐Ÿง  torch.no_grad()

During inference, gradients are generally unnecessary.

with torch.no_grad():

    output = model(
        x
    )

This reduces unnecessary autograd tracking and can lower memory usage.


๐Ÿง  Training vs Inference

flowchart LR

    TRAIN["Training"]

    TRAIN --> MODE1["model.train()"]
    MODE1 --> FORWARD1["Forward"]
    FORWARD1 --> LOSS["Loss"]
    LOSS --> BACK["Backward"]
    BACK --> UPDATE["Optimizer Step"]

    INFER["Inference"]

    INFER --> MODE2["model.eval()"]
    MODE2 --> NOGRAD["torch.no_grad()"]
    NOGRAD --> FORWARD2["Forward"]
    FORWARD2 --> OUTPUT["Prediction"]

๐Ÿง  PyTorch Training Workflow

The fundamental training loop is:

for x_batch, y_batch in train_loader:

    optimizer.zero_grad()

    predictions = model(
        x_batch
    )

    loss = loss_fn(
        predictions,
        y_batch
    )

    loss.backward()

    optimizer.step()

The workflow is:

Zero Gradients
      โ†“
Forward Pass
      โ†“
Calculate Loss
      โ†“
Backward Pass
      โ†“
Update Parameters

๐Ÿง  PyTorch Training Loop Architecture

flowchart TD

    DATA["DataLoader"]

    BATCH["Batch"]

    ZERO["optimizer.zero_grad()"]

    FORWARD["model(x)"]

    LOSS["loss_fn()"]

    BACKWARD["loss.backward()"]

    STEP["optimizer.step()"]

    DATA --> BATCH
    BATCH --> ZERO
    ZERO --> FORWARD
    FORWARD --> LOSS
    LOSS --> BACKWARD
    BACKWARD --> STEP
    STEP --> BATCH

๐Ÿง  Loss Functions

PyTorch provides many loss functions.

Common examples:

nn.MSELoss()
nn.L1Loss()
nn.CrossEntropyLoss()
nn.BCEWithLogitsLoss()

Example:

loss_fn = nn.CrossEntropyLoss()

๐Ÿง  Optimizers

PyTorch provides:

torch.optim

Common optimizers include:

SGD
Adam
AdamW
RMSprop

Example:

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

๐Ÿง  Optimizer Workflow

flowchart LR

    MODEL["Model"]

    PRED["Prediction"]

    LOSS["Loss"]

    GRAD["Gradients"]

    OPT["Optimizer"]

    UPDATE["Updated Parameters"]

    MODEL --> PRED
    PRED --> LOSS
    LOSS --> GRAD
    GRAD --> OPT
    OPT --> UPDATE
    UPDATE --> MODEL

๐Ÿงช Complete PyTorch Classification Example

import torch
import torch.nn as nn


device = torch.device(
    "cuda"
    if torch.cuda.is_available()
    else "cpu"
)


class Classifier(
    nn.Module
):

    def __init__(self):

        super().__init__()

        self.network = nn.Sequential(

            nn.Linear(
                784,
                128
            ),

            nn.ReLU(),

            nn.Linear(
                128,
                64
            ),

            nn.ReLU(),

            nn.Linear(
                64,
                10
            )
        )

    def forward(
        self,
        x
    ):

        return self.network(
            x
        )


model = Classifier().to(
    device
)


loss_fn = nn.CrossEntropyLoss()


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

Training:

for epoch in range(
    10
):

    model.train()

    for x_batch, y_batch in train_loader:

        x_batch = x_batch.to(
            device
        )

        y_batch = y_batch.to(
            device
        )

        optimizer.zero_grad()

        predictions = model(
            x_batch
        )

        loss = loss_fn(
            predictions,
            y_batch
        )

        loss.backward()

        optimizer.step()

๐Ÿง  Why CrossEntropyLoss Expects Logits

For multi-class classification, the model commonly returns raw logits:

outputs = model(
    x
)

The final layer is usually:

nn.Linear(
    64,
    10
)

without explicitly applying:

Softmax

when using:

nn.CrossEntropyLoss()

CrossEntropyLoss internally combines the required log-softmax and negative log-likelihood behavior.

Therefore:

Model
 โ†“
Raw Logits
 โ†“
CrossEntropyLoss

is the common PyTorch pattern.


๐Ÿง  PyTorch Classification Architecture

flowchart LR

    INPUT["Input"]

    FC1["Linear"]

    RELU1["ReLU"]

    FC2["Linear"]

    RELU2["ReLU"]

    LOGITS["Class Logits"]

    LOSS["CrossEntropyLoss"]

    INPUT --> FC1
    FC1 --> RELU1
    RELU1 --> FC2
    FC2 --> RELU2
    RELU2 --> LOGITS
    LOGITS --> LOSS

๐Ÿงช PyTorch Regression Example

class RegressionModel(
    nn.Module
):

    def __init__(self):

        super().__init__()

        self.network = nn.Sequential(

            nn.Linear(
                10,
                64
            ),

            nn.ReLU(),

            nn.Linear(
                64,
                32
            ),

            nn.ReLU(),

            nn.Linear(
                32,
                1
            )
        )

    def forward(
        self,
        x
    ):

        return self.network(
            x
        )

Loss:

loss_fn = nn.MSELoss()

Optimizer:

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

๐Ÿง  Classification vs Regression

Problem Output Typical Loss
Binary Classification Logit / Sigmoid interpretation BCEWithLogitsLoss
Multi-Class Classification Class Logits CrossEntropyLoss
Regression Continuous Value MSELoss / L1Loss

๐Ÿง  PyTorch Parameters

Parameters are represented by:

nn.Parameter

They are tensors that PyTorch tracks as trainable model parameters.

Example:

self.weight = nn.Parameter(
    torch.randn(
        10,
        5
    )
)

๐Ÿง  Parameter Registration

When layers are assigned as attributes of an nn.Module:

self.fc = nn.Linear(
    10,
    5
)

PyTorch automatically registers the layer and its parameters.

This allows:

model.parameters()

to discover them.


๐Ÿง  state_dict()

A PyTorch model's parameters and persistent buffers can be accessed through:

model.state_dict()

Example:

state = model.state_dict()

for key, value in state.items():

    print(
        key,
        value.shape
    )

๐Ÿ’พ Saving Model Weights

A common approach is:

torch.save(
    model.state_dict(),
    "model.pth"
)

Load:

model.load_state_dict(
    torch.load(
        "model.pth",
        weights_only=True
    )
)

The exact loading approach can depend on the PyTorch version and whether the artifact contains weights, a checkpoint, or another object.


๐Ÿง  Model Checkpoint

A production checkpoint may contain more than model weights.

For example:

checkpoint = {

    "model_state_dict":
        model.state_dict(),

    "optimizer_state_dict":
        optimizer.state_dict(),

    "epoch":
        epoch,

    "loss":
        loss
}

Save:

torch.save(
    checkpoint,
    "checkpoint.pth"
)

This allows training to resume more completely.


๐Ÿง  Checkpoint Architecture

flowchart TD

    TRAIN["Training"]

    MODEL["Model State"]

    OPT["Optimizer State"]

    META["Training Metadata"]

    CHECK["Checkpoint"]

    TRAIN --> MODEL
    TRAIN --> OPT
    TRAIN --> META

    MODEL --> CHECK
    OPT --> CHECK
    META --> CHECK

    CHECK --> RESUME["Resume Training"]

๐Ÿง  PyTorch Device + Model

A common production pattern:

device = torch.device(
    "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

model = model.to(
    device
)

For each batch:

x_batch = x_batch.to(
    device
)

y_batch = y_batch.to(
    device
)

๐Ÿง  GPU Training Architecture

flowchart LR

    DATA["CPU Dataset"]

    LOADER["DataLoader"]

    CPU_BATCH["CPU Batch"]

    GPU["GPU"]

    MODEL["PyTorch Model"]

    GRAD["Gradients"]

    UPDATE["Parameter Update"]

    DATA --> LOADER
    LOADER --> CPU_BATCH
    CPU_BATCH --> GPU
    GPU --> MODEL
    MODEL --> GRAD
    GRAD --> UPDATE
    UPDATE --> MODEL

โšก CUDA Device Selection

For multiple GPUs:

device = torch.device(
    "cuda:0"
)

Check GPU count:

torch.cuda.device_count()

Get device name:

torch.cuda.get_device_name(
    0
)

๐Ÿง  CUDA Memory

Useful information:

torch.cuda.memory_allocated()

and:

torch.cuda.memory_reserved()

These help diagnose GPU memory consumption.


๐Ÿง  CPU โ†” GPU Data Movement

Moving data between devices has a cost.

CPU Memory
    โ”‚
    โ”‚ Data Transfer
    โ–ผ
GPU Memory

Therefore, efficient Deep Learning systems try to avoid unnecessary transfers.


โš  Common Device Mistake

Avoid repeatedly moving tensors:

for batch in data:

    x = x.to(
        device
    )

    # unnecessary repeated transfers

Instead, structure the pipeline so that device movement occurs at a predictable point.


๐Ÿง  Tensor Memory Sharing

When converting data from NumPy:

x = torch.from_numpy(
    numpy_array
)

PyTorch may share the underlying memory with the NumPy array.

This differs from:

torch.tensor(
    numpy_array
)

which generally creates a new tensor and copies the data.

This distinction can matter for performance and mutation behavior.


๐Ÿง  Detaching Tensors

If a tensor is part of an autograd graph:

x.detach()

creates a tensor that does not require gradient tracking through that history.

Example:

prediction = model(
    x
)

prediction_cpu = (
    prediction
    .detach()
    .cpu()
)

๐Ÿง  NumPy Conversion

For tensors that do not require gradients:

array = tensor.numpy()

For tensors that require gradients:

array = (
    tensor
    .detach()
    .cpu()
    .numpy()
)

๐Ÿง  detach() Mental Model

flowchart LR

    TENSOR["Tensor"]

    GRAPH["Autograd Graph"]

    DETACH["detach()"]

    OUTPUT["Detached Tensor"]

    TENSOR --> GRAPH
    GRAPH --> DETACH
    DETACH --> OUTPUT

๐Ÿง  torch.no_grad() vs detach()

These are related but serve different purposes.

torch.no_grad()

Disables gradient tracking for operations inside the context.

with torch.no_grad():

    output = model(
        x
    )

detach()

Creates a tensor disconnected from the current autograd graph.

output = prediction.detach()

๐Ÿง  inference_mode()

For inference workloads, PyTorch also provides:

with torch.inference_mode():

    output = model(
        x
    )

This can provide additional performance benefits compared with ordinary no_grad() in appropriate inference scenarios.


๐Ÿง  PyTorch Model Lifecycle

A typical workflow is:

flowchart LR

    DATA["Dataset"]

    TENSOR["Tensor"]

    MODEL["nn.Module"]

    COMPILE["Configure Loss + Optimizer"]

    TRAIN["Training"]

    CHECKPOINT["Checkpoint"]

    EVAL["Evaluation"]

    DEPLOY["Deployment"]

    DATA --> TENSOR
    TENSOR --> MODEL
    MODEL --> COMPILE
    COMPILE --> TRAIN
    TRAIN --> CHECKPOINT
    CHECKPOINT --> EVAL
    EVAL --> DEPLOY

๐Ÿง  PyTorch vs Keras Mental Mapping

Concept Keras / TensorFlow PyTorch
Tensor tf.Tensor torch.Tensor
Layer tf.keras.layers.Layer nn.Module
Model tf.keras.Model nn.Module
Dense Dense nn.Linear
Conv2D Conv2D nn.Conv2d
ReLU ReLU nn.ReLU
Gradient GradientTape Autograd
Optimizer tf.keras.optimizers torch.optim
Dataset Pipeline tf.data Dataset / DataLoader
Training model.fit() Training loop
Evaluation Mode training=False model.eval()
No Gradients Context / inference logic torch.no_grad() / inference_mode()
Save Model .keras / SavedModel workflows state_dict() / checkpoints

๐Ÿง  The Core PyTorch Training Equation

The fundamental parameter update is:

[ \theta_{t+1} = \theta_t - \eta \nabla_\theta L ]

Where:

ฮธ  = Model Parameters
ฮท  = Learning Rate
L  = Loss
โˆ‡ฮธL = Gradient

The PyTorch training loop implements this concept through:

loss.backward()
        โ†“
Gradients
        โ†“
optimizer.step()
        โ†“
Updated Parameters

๐Ÿข Enterprise Perspective

PyTorch should not be viewed only as a model-building library.

In production, it is one component of a larger Deep Learning platform:

Data Sources
     โ†“
Data Pipeline
     โ†“
Dataset / DataLoader
     โ†“
PyTorch Model
     โ†“
Training Infrastructure
     โ†“
GPU / Accelerator
     โ†“
Checkpoint
     โ†“
Model Validation
     โ†“
Model Registry
     โ†“
Serving
     โ†“
Monitoring

Production concerns include:

  • Dataset versioning
  • Reproducibility
  • Configuration management
  • GPU utilization
  • Checkpointing
  • Model versioning
  • Experiment tracking
  • Model validation
  • Deployment
  • Monitoring
  • Security
  • Cost management

Production Insight

PyTorch gives you significant control over the Deep Learning execution model.

That flexibility is powerful, but it also means the engineering team must explicitly manage:

Device Placement
+
Memory
+
Gradient Lifecycle
+
Checkpointing
+
Training State
+
Reproducibility

A model that trains successfully on a developer laptop is not automatically production-ready.


โš  Common Mistakes

Avoid these common PyTorch mistakes:

  • Mixing CPU and GPU tensors
  • Forgetting to move the model to the target device
  • Forgetting to move input batches to the target device
  • Forgetting optimizer.zero_grad()
  • Calling backward() without understanding gradient accumulation
  • Forgetting model.train() during training
  • Forgetting model.eval() during evaluation
  • Calculating inference with unnecessary gradient tracking
  • Using an inappropriate output activation with a chosen loss function
  • Applying Softmax before CrossEntropyLoss unnecessarily
  • Ignoring tensor shape conventions
  • Confusing view() and reshape()
  • Misusing squeeze() and accidentally removing meaningful dimensions
  • Using incorrect permute() ordering
  • Performing unnecessary CPU/GPU transfers
  • Converting GPU tensors directly to NumPy
  • Failing to detach tensors before converting them for logging
  • Saving only partial training state when resuming is required
  • Ignoring GPU memory consumption
  • Using unnecessarily large batch sizes
  • Creating tensors on the wrong device inside the model

๐Ÿง  Interview Questions

Beginner

1. What is PyTorch?

PyTorch is a Deep Learning framework providing tensor computation, automatic differentiation, neural network abstractions, optimization, and hardware acceleration.

2. What is a tensor?

A tensor is a multidimensional numerical data structure used as the fundamental data representation in PyTorch.

3. What is nn.Module?

nn.Module is the base class used to define neural network models and reusable neural network components in PyTorch.

4. What does forward() do?

It defines how input data flows through a PyTorch model.

5. What is requires_grad=True?

It tells PyTorch to track operations involving the tensor so gradients can be computed through autograd.


Intermediate

6. What is autograd?

PyTorch's automatic differentiation system that records differentiable operations and computes gradients during backward propagation.

7. What is the standard PyTorch training loop?

optimizer.zero_grad()
        โ†“
forward
        โ†“
loss
        โ†“
loss.backward()
        โ†“
optimizer.step()

8. Why call optimizer.zero_grad()?

Because PyTorch gradients accumulate by default. Existing gradients need to be cleared before computing the next update.

9. What is the difference between model.train() and model.eval()?

They switch the model between training and evaluation behavior, which affects layers such as Dropout and Batch Normalization.

10. Why use torch.no_grad() during inference?

It disables gradient tracking for the enclosed operations, reducing unnecessary computation and memory usage.

11. What is state_dict()?

It provides the model's parameters and persistent buffers in a dictionary-like structure that is commonly used for saving and loading model state.

12. Why use DataLoader?

DataLoader provides batching, iteration, shuffling, and other mechanisms for efficiently feeding data into a training loop. Its details are covered in the next chapter.


Advanced

13. Why does PyTorch use nn.Module for both layers and models?

Because complex neural networks can be composed hierarchically from reusable modules. This allows PyTorch to recursively track parameters and submodules.

14. Why are model and tensors required to be on compatible devices?

Operations generally require tensors participating in the same computation to reside on compatible devices.

15. Why does CrossEntropyLoss typically receive raw logits?

PyTorch's CrossEntropyLoss combines the relevant log-softmax and negative-log-likelihood computation internally, so an explicit Softmax layer is normally unnecessary before it.

16. What is the difference between detach() and torch.no_grad()?

detach() disconnects a tensor from its existing autograd history. torch.no_grad() disables gradient tracking for operations executed inside its context.

17. Why can view() fail when reshape() works?

view() requires a compatible memory layout, while reshape() can create a copy when necessary.

18. Why is permute() important in Computer Vision?

Different frameworks and operations expect different dimension orders. PyTorch vision models commonly use:

N ร— C ร— H ร— W

so tensors may need to be permuted into that format.

19. What consumes GPU memory during training?

Typically:

Parameters
+
Gradients
+
Optimizer State
+
Activations
+
Input Batches

20. How would you design production PyTorch training?

Separate:

Dataset
DataLoader
Model
Loss
Optimizer
Training Loop
Evaluation
Checkpointing
Configuration
Logging

and make device placement, reproducibility, checkpointing, and monitoring explicit.


๐Ÿงช Practical Exercises

Exercise 1 โ€” Tensor Fundamentals

Create:

Scalar
Vector
Matrix
3D Tensor
4D Tensor

For each tensor print:

Value
Shape
Rank
Data Type
Device

Exercise 2 โ€” Tensor Operations

Implement:

Addition
Subtraction
Multiplication
Matrix Multiplication
Reshape
Transpose
Permute
Squeeze
Unsqueeze
Mean
Sum

Verify the resulting shapes.


Exercise 3 โ€” Autograd

Create:

x = torch.tensor(
    3.0,
    requires_grad=True
)

Calculate:

[ y=x3+2x2+x ]

and use:

y.backward()

to calculate the gradient.


Exercise 4 โ€” Build a Classification Model

Create:

Input
 โ†“
Linear
 โ†“
ReLU
 โ†“
Linear
 โ†“
ReLU
 โ†“
Linear
 โ†“
Class Logits

Use:

nn.CrossEntropyLoss()

for training.


Exercise 5 โ€” Build a Regression Model

Create:

Input
 โ†“
Linear
 โ†“
ReLU
 โ†“
Linear
 โ†“
Regression Output

Use:

nn.MSELoss()

and:

AdamW

for optimization.


Exercise 6 โ€” CPU vs GPU

Detect:

torch.cuda.is_available()

Train the same model using:

CPU
GPU

Compare:

Training Time
Throughput
Memory Usage

Exercise 7 โ€” Model Checkpointing

Save:

Model State
Optimizer State
Epoch
Training Loss

Create a checkpoint and resume training from it.


Exercise 8 โ€” Tensor Layout

Create an image batch in:

N ร— H ร— W ร— C

and convert it to:

N ร— C ร— H ร— W

using:

permute()

Verify the resulting shape.


๐Ÿ“Œ Key Takeaways

  • PyTorch is a tensor-based Deep Learning framework.
  • Tensors are the fundamental data structure in PyTorch.
  • Tensor rank represents the number of dimensions.
  • Tensor shape describes the size of each dimension.
  • Tensor data types affect memory, precision, and performance.
  • PyTorch supports CPU and accelerator-based tensor computation.
  • Device management is critical when using GPUs.
  • torch.autograd provides automatic differentiation.
  • requires_grad=True enables gradient tracking for tensors.
  • backward() computes gradients through the autograd graph.
  • PyTorch gradients accumulate by default.
  • optimizer.zero_grad() clears previous gradients.
  • optimizer.step() updates model parameters.
  • nn.Module is the fundamental abstraction for PyTorch neural networks.
  • forward() defines the model's forward computation.
  • nn.Linear represents a fully connected layer.
  • nn.Sequential is useful for simple linear architectures.
  • Complex architectures should generally use custom nn.Module implementations.
  • model.train() enables training behavior.
  • model.eval() enables evaluation behavior.
  • torch.no_grad() and torch.inference_mode() help avoid unnecessary gradient tracking during inference.
  • state_dict() is commonly used for model state serialization.
  • Efficient GPU training requires careful management of device placement and memory.
  • PyTorch gives developers significant control over the training process.
  • That flexibility also creates greater responsibility for training infrastructure, reproducibility, checkpointing, and production reliability.

๐Ÿ“š Further Reading

Continue with:

The next chapter goes deeper into PyTorch Autograd, Dataset, and DataLoader, connecting tensor computation with efficient real-world data pipelines and training workflows.


โžก๏ธ Next Chapter

17. PyTorch Autograd, Dataset and DataLoader


Enterprise AI Engineering Handbook
Building Production-Grade Enterprise AI Systems โ€” One Chapter at a Time.