∇GradientSchool
Intermediate
30 min read
#PyTorch#Tensors#Autograd#nn.Module#CUDA#Training Loops#Frameworks

26. Deep Learning Frameworks (PyTorch Core Architecture)

Mastering the PyTorch engine: tensor memory layouts, autograd computational graphs, custom nn.Module building blocks, dataset/DataLoader pipelining, and multi-GPU training.

Deep Learning Frameworks: Complete Notes (Beginner to Advanced)

Introduction#

Deep Learning Frameworks are software libraries that provide the tools required to build, train, evaluate, save, and deploy neural networks.

Two major frameworks are:

›PyTorch
TensorFlow

Both provide:

  • Tensor operations
  • Automatic differentiation
  • Neural-network building blocks
  • Loss functions
  • Optimizers
  • GPU acceleration
  • Model training utilities
  • Model saving and loading

A useful high-level view is:

text
                 Deep Learning Frameworks
                          │
             ┌────────────┴────────────┐
             │                         │
          PyTorch                  TensorFlow
             │                         │
       ┌─────┼─────┐             ┌─────┼─────┐
       │     │     │             │     │     │
    Tensors Autograd nn.Module  Tensors Keras Layers
       │     │     │             │     │     │
       └─────┴─────┘             └─────┴─────┘
             │                         │
        Training Loop              Training
        DataLoader                 Models
        Optimizers                 GPU Training

This topic focuses specifically on the requested PyTorch and TensorFlow/Keras components.

1. PyTorch

1.1 What is PyTorch?#

PyTorch is an open-source deep learning framework that provides tensor computation, automatic differentiation, neural-network modules, optimizers, data-loading utilities, and hardware acceleration.

PyTorch is commonly used for:

  • Deep learning research
  • Computer vision
  • Natural language processing
  • Generative AI
  • Model training
  • Model experimentation
  • Production inference

The basic workflow is:

text
Data
 ↓
Tensor
 ↓
Model
 ↓
Prediction
 ↓
Loss
 ↓
Autograd
 ↓
Gradients
 ↓
Optimizer
 ↓
Updated Model

2. PyTorch Tensors

2.1 What is a Tensor?#

A tensor is a multidimensional numerical array.

You can think of tensors as a generalization of:

text
Scalar → 0 dimensions
Vector → 1 dimension
Matrix → 2 dimensions
Tensor → 3+ dimensions

Examples:

text
Scalar:
5

Vector:
[1, 2, 3]

Matrix:
[[1, 2],
 [3, 4]]

3D Tensor:
multiple matrices stacked together

Deep learning models perform most computations using tensors.

2.2 Creating Tensors#

🐍 Python
import torch

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

print(x)
print(x.shape)
print(x.dtype)

Output conceptually:

text
tensor([1, 2, 3])
shape: torch.Size([3])
dtype: torch.int64

Floating-point tensors are commonly used for neural-network computations:

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

2.3 Common Tensor Creation Functions#

🐍 Python
torch.zeros(3, 4)
torch.ones(3, 4)
torch.randn(3, 4)
torch.rand(3, 4)
torch.arange(10)

Examples:

🐍 Python
x = torch.zeros(2, 3)
y = torch.ones(2, 3)
z = torch.randn(2, 3)

2.4 Tensor Shape#

Shape describes the size of every dimension.

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

print(x.shape)

A common image-batch interpretation is:

text
32    → batch size
3     → channels
224   → height
224   → width

Therefore:

›(batch, channels, height, width)

2.5 Tensor Operations#

🐍 Python
a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])

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

Matrix multiplication:

🐍 Python
A = torch.randn(3, 4)
B = torch.randn(4, 2)

C = A @ B

or:

🐍 Python
C = torch.matmul(A, B)

2.6 Reshaping Tensors#

🐍 Python
x = torch.arange(12)

x = x.reshape(3, 4)

The number of elements must remain the same.

Mathematical Formulation
12 elements

3 × 4 = 12

Another common operation:

🐍 Python
x = x.view(3, 4)

reshape() is generally more flexible, while view() has stricter memory-layout requirements.

2.7 Device Placement#

A tensor can reside on:

›CPU
GPU

Example:

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

x = torch.randn(3, 4).to(device)

The model and tensors involved in the same operation generally need to be on compatible devices.

3. Autograd

3.1 What is Autograd?#

Autograd is PyTorch's automatic differentiation system.

It automatically computes gradients needed for optimization.

Instead of manually deriving:

›∂Loss / ∂Weight

PyTorch can calculate it automatically.

3.2 requires_grad#

If a tensor should be differentiated with respect to during training:

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

Now PyTorch tracks operations involving x.

Example:

🐍 Python
y = x ** 2

Since:

Mathematical Formulation
y = x²

the derivative is:

Mathematical Formulation
dy/dx = 2x

At:

Mathematical Formulation
x = 2

the gradient is:

›4

PyTorch:

🐍 Python
y.backward()

print(x.grad)

Output:

›tensor(4.)

3.3 Computational Graph#

Autograd builds a computational graph as operations are performed.

For:

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

y = x ** 2
z = y + 3

z.backward()

Conceptually:

text
x
│
▼
x²
│
▼
y
│
+ 3
│
▼
z

Calling:

🐍 Python
z.backward()

computes gradients through the graph.

3.4 backward()#

The .backward() method performs backpropagation from a scalar output.

Example:

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

y = x ** 2
y.backward()

print(x.grad)

Result:

›6

because:

Mathematical Formulation
dy/dx = 2x

and:

Mathematical Formulation
2(3) = 6

3.5 Gradient Accumulation#

PyTorch gradients accumulate by default.

Therefore, training loops commonly contain:

🐍 Python
optimizer.zero_grad()

before:

🐍 Python
loss.backward()

Typical sequence:

text
zero gradients
      ↓
forward pass
      ↓
calculate loss
      ↓
backward pass
      ↓
optimizer step

4. nn.Module

4.1 What is nn.Module?#

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

A model usually inherits from:

🐍 Python
nn.Module

Example:

🐍 Python
import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.fc1 = nn.Linear(10, 32)
        self.fc2 = nn.Linear(32, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)

        return x

4.2 init()#

The __init__() method defines the model's layers and parameters.

🐍 Python
self.fc1 = nn.Linear(10, 32)

creates a trainable linear layer.

4.3 forward()#

The forward() method defines how data moves through the model.

🐍 Python
def forward(self, x):
    x = self.fc1(x)
    x = torch.relu(x)
    return self.fc2(x)

When you write:

🐍 Python
output = model(x)

PyTorch calls the model's forward() method.

4.4 Parameters#

Trainable parameters can be accessed with:

🐍 Python
model.parameters()

For example:

🐍 Python
for parameter in model.parameters():
    print(parameter.shape)

This is what optimizers use to update the model.

5. Dataset

5.1 What is Dataset?#

torch.utils.data.Dataset represents a collection of training or evaluation examples.

A custom dataset normally implements:

›__len__()
__getitem__()

Example:

🐍 Python
from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, X, y):
        self.X = X
        self.y = y

    def __len__(self):
        return len(self.X)

    def __getitem__(self, index):
        return self.X[index], self.y[index]

5.2 len()#

Returns the number of samples.

🐍 Python
def __len__(self):
    return len(self.X)

5.3 getitem()#

Returns one sample.

🐍 Python
def __getitem__(self, index):
    return self.X[index], self.y[index]

A returned sample might be:

›(features, label)

6. DataLoader

6.1 What is DataLoader?#

DataLoader takes a Dataset and provides an iterable over batches.

Instead of processing:

›1 sample at a time

it can provide:

›batch of samples

Example:

🐍 Python
from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True
)

6.2 Important DataLoader Arguments#

batch_size#

Controls the number of samples in each batch.

🐍 Python
batch_size=32

means approximately:

›32 samples per batch

except possibly the final batch.

shuffle#

🐍 Python
shuffle=True

randomizes the order of samples at the beginning of each epoch.

This is commonly used for training.

num_workers#

Controls worker processes used to load data.

🐍 Python
DataLoader(
    dataset,
    batch_size=32,
    num_workers=4
)

The useful value depends on the system and data-loading workload.

6.3 Dataset vs DataLoader#

ComponentPurpose
DatasetDefines how individual samples are accessed
DataLoaderOrganizes samples into batches and provides iteration

Conceptually:

text
Dataset
[1, 2, 3, 4, 5, 6, ...]
          ↓
      DataLoader
          ↓
[1,2] [3,4] [5,6]

7. Training Loop

7.1 What is a Training Loop?#

A training loop repeatedly:

  1. Gets a batch of data
  2. Performs a forward pass
  3. Calculates loss
  4. Computes gradients
  5. Updates parameters

The core flow is:

text
Batch
 ↓
Forward Pass
 ↓
Prediction
 ↓
Loss
 ↓
Backward Pass
 ↓
Gradients
 ↓
Optimizer Step
 ↓
Updated Model

7.2 Basic Training Loop#

🐍 Python
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()

7.3 Why zero_grad()?#

Gradients accumulate by default.

Therefore:

🐍 Python
optimizer.zero_grad()

clears gradients from the previous optimization step.

7.4 Why loss.backward()?#

This calculates gradients of the loss with respect to trainable parameters.

🐍 Python
loss.backward()

Conceptually:

text
Loss
 ↓
Autograd
 ↓
∂Loss/∂W
∂Loss/∂b
...

7.5 Why optimizer.step()?#

The optimizer uses the calculated gradients to update model parameters.

🐍 Python
optimizer.step()

For gradient descent, conceptually:

Mathematical Formulation
parameter_new
=
parameter_old
-
learning_rate × gradient

7.6 Epoch#

An epoch means one complete pass through the training dataset.

Example:

🐍 Python
for epoch in range(10):

    for X_batch, y_batch in train_loader:
        ...

This trains for:

›10 epochs

8. Validation Loop

8.1 What is Validation?#

Validation evaluates the model on data that is not used to update its parameters.

The purpose is to estimate how well the model generalizes during training.

Typical workflow:

text
Training Data
     ↓
Update Model

Validation Data
     ↓
Evaluate Model
     ↓
No Parameter Updates

8.2 model.eval()#

Before validation:

🐍 Python
model.eval()

This switches certain layers into evaluation behavior.

For example:

  • Dropout stops randomly dropping activations
  • Batch Normalization uses stored running statistics

8.3 torch.no_grad()#

During validation:

🐍 Python
with torch.no_grad():

disables gradient tracking for the operations inside the block.

Example:

🐍 Python
model.eval()

with torch.no_grad():

    for X_batch, y_batch in val_loader:

        predictions = model(X_batch)

        loss = loss_fn(predictions, y_batch)

This reduces memory usage and avoids unnecessary gradient computation.

8.4 Returning to Training Mode#

After validation:

🐍 Python
model.train()

restores training behavior.

8.5 Complete Training + Validation Flow#

🐍 Python
for epoch in range(num_epochs):

    # Training
    model.train()

    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()

    # Validation
    model.eval()

    with torch.no_grad():

        for X_batch, y_batch in val_loader:

            predictions = model(X_batch)

            val_loss = loss_fn(
                predictions,
                y_batch
            )

9. Optimizers

9.1 What is an Optimizer?#

An optimizer updates trainable model parameters using their gradients.

General idea:

text
Loss
 ↓
Gradients
 ↓
Optimizer
 ↓
Updated Parameters

9.2 SGD#

Stochastic Gradient Descent (SGD) uses gradients to update parameters.

Conceptually:

Mathematical Formulation
θ_new = θ_old - η∇θL

where:

  • θ = model parameters
  • η = learning rate
  • L = loss

PyTorch:

🐍 Python
optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.01
)

9.3 Momentum#

Momentum uses information from previous gradients to influence the current update.

Conceptually:

text
current update
+
previous update information

PyTorch:

🐍 Python
optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9
)

9.4 Adam#

Adam (Adaptive Moment Estimation) maintains moving estimates related to the first and second moments of gradients.

A simplified conceptual view is:

text
gradient
   ↓
estimate first moment
estimate second moment
   ↓
adaptive parameter update

PyTorch:

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

Adam is widely used because it often provides convenient optimization for many neural-network problems.

9.5 AdamW#

AdamW separates weight decay from the gradient-based Adam update.

🐍 Python
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=0.001,
    weight_decay=0.01
)

AdamW is commonly used when training modern deep-learning models.

9.6 Optimizer Comparison#

OptimizerMain Idea
SGDGradient-based updates
SGD + MomentumAdds accumulated update direction
AdamAdaptive first/second-moment-based updates
AdamWAdam with decoupled weight decay

10. Loss Functions

10.1 What is a Loss Function?#

A loss function measures how different the model's prediction is from the target.

text
Prediction
     +
Target
     ↓
Loss

The optimizer tries to minimize the loss.

10.2 Mean Squared Error#

MSE is commonly used for regression.

Mathematical Formulation
MSE = (1/n) Σ(y_pred - y_true)²

PyTorch:

🐍 Python
loss_fn = nn.MSELoss()

Example:

🐍 Python
loss = loss_fn(predictions, targets)

10.3 Cross-Entropy Loss#

Cross-entropy is commonly used for classification.

PyTorch:

🐍 Python
loss_fn = nn.CrossEntropyLoss()

For standard multi-class classification, the model typically produces raw logits:

›[logit_class_1, logit_class_2, ...]

CrossEntropyLoss internally combines the relevant log-softmax operation with negative log-likelihood.

Therefore, you normally should not apply softmax before passing logits to CrossEntropyLoss.

10.4 Binary Cross-Entropy#

For binary classification, a common choice is:

🐍 Python
nn.BCEWithLogitsLoss()

It combines sigmoid behavior with binary cross-entropy in a numerically stable way.

Therefore, the model should normally output logits rather than applying sigmoid before the loss.

10.5 Loss Function Examples#

TaskCommon Loss
RegressionMSELoss
Multi-class classificationCrossEntropyLoss
Binary classificationBCEWithLogitsLoss
Multi-label classificationBCEWithLogitsLoss

The appropriate loss depends on the task and target representation.

11. GPU Training

11.1 Why Use a GPU?#

Neural networks perform large numbers of parallel numerical operations.

GPUs are designed to perform many such operations efficiently.

For suitable workloads:

›CPU → slower training
GPU → faster training

The actual speedup depends on the model, batch size, hardware, and workload.

11.2 Selecting a Device#

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

Move model:

🐍 Python
model = model.to(device)

Move tensors:

🐍 Python
X = X.to(device)
y = y.to(device)

11.3 GPU Training Loop#

🐍 Python
model.to(device)

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()

The important requirement is that the model and tensors involved in computation are on compatible devices.

12. Model Saving

12.1 Saving Model Parameters#

A common PyTorch practice is to save the model's state_dict.

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

A state_dict contains the model's parameter and persistent buffer values.

12.2 Why Save state_dict?#

Saving the parameter state is generally more flexible than serializing the entire model object.

The model architecture can be recreated in code and the learned parameters loaded into it.

text
Model Architecture
       +
state_dict
       ↓
Reconstructed Model

13. Model Loading

13.1 Loading state_dict#

First recreate the architecture:

🐍 Python
model = SimpleModel()

Then load the parameters:

🐍 Python
model.load_state_dict(
    torch.load("model.pth")
)

For inference:

🐍 Python
model.eval()

13.2 Device-Aware Loading#

If the checkpoint may be loaded on a different device:

🐍 Python
state_dict = torch.load(
    "model.pth",
    map_location=device
)

model.load_state_dict(state_dict)

14. Checkpoints

14.1 What is a Checkpoint?#

A checkpoint is a saved snapshot of training state.

It can contain:

text
Model parameters
+
Optimizer state
+
Epoch
+
Validation loss
+
Other training information

Example:

🐍 Python
checkpoint = {
    "epoch": epoch,
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "loss": loss.item()
}

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

14.2 Loading a Checkpoint#

🐍 Python
checkpoint = torch.load(
    "checkpoint.pth",
    map_location=device
)

model.load_state_dict(
    checkpoint["model_state_dict"]
)

optimizer.load_state_dict(
    checkpoint["optimizer_state_dict"]
)

epoch = checkpoint["epoch"]

This allows training to continue from a saved point.

14.3 Model vs Checkpoint#

text
Model file
→ primarily stores learned model state

Checkpoint
→ stores model state + training state

The exact contents are implementation-dependent.

15. TensorFlow

15.1 What is TensorFlow?#

TensorFlow is an open-source machine-learning framework developed by Google.

It provides:

  • Tensor operations
  • Automatic differentiation
  • Neural-network APIs
  • Optimizers
  • Loss functions
  • GPU acceleration
  • Model training and serialization

TensorFlow can be used directly through its lower-level APIs or through Keras, its high-level neural-network API integrated into TensorFlow.

16. TensorFlow Tensors

16.1 What is a TensorFlow Tensor?#

A TensorFlow tensor is a multidimensional array used for numerical computation.

Example:

🐍 Python
import tensorflow as tf

x = tf.constant([1.0, 2.0, 3.0])

print(x)
print(x.shape)
print(x.dtype)

16.2 Creating Tensors#

🐍 Python
tf.zeros((2, 3))
tf.ones((2, 3))
tf.random.normal((2, 3))
tf.random.uniform((2, 3))
tf.range(10)

Example:

🐍 Python
x = tf.random.normal((32, 10))

16.3 Tensor Operations#

🐍 Python
a = tf.constant([1.0, 2.0, 3.0])
b = tf.constant([4.0, 5.0, 6.0])

print(a + b)
print(a * b)

Matrix multiplication:

🐍 Python
A = tf.random.normal((3, 4))
B = tf.random.normal((4, 2))

C = tf.matmul(A, B)

16.4 Tensor Reshaping#

🐍 Python
x = tf.range(12)

x = tf.reshape(x, (3, 4))

The number of elements must remain unchanged.

17. Keras

17.1 What is Keras?#

Keras is a high-level deep-learning API used with TensorFlow.

It provides a convenient way to define:

  • Layers
  • Models
  • Loss functions
  • Optimizers
  • Training workflows

Instead of manually implementing many low-level training details, Keras provides higher-level abstractions.

A common TensorFlow workflow is:

text
TensorFlow
   │
   ▼
 Keras
   │
   ├── Layers
   ├── Models
   ├── Losses
   ├── Optimizers
   └── Training

18. Keras Layers

18.1 What is a Layer?#

A layer transforms its input into an output.

Examples:

🐍 Python
tf.keras.layers.Dense(64)
tf.keras.layers.Conv2D(32, 3)
tf.keras.layers.Dropout(0.2)
tf.keras.layers.BatchNormalization()

A dense layer:

🐍 Python
layer = tf.keras.layers.Dense(32)

contains trainable weights that are created based on the input shape when the layer is built.

18.2 Sequential Model#

For a simple linear stack of layers:

🐍 Python
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation="relu"),
    tf.keras.layers.Dense(10)
])

Conceptually:

text
Input
  ↓
Dense 64
  ↓
ReLU
  ↓
Dense 10
  ↓
Output

19. Keras Models

19.1 What is a Model?#

A Keras model represents a neural network and defines how inputs are transformed into outputs.

Common approaches include:

text
Sequential API
Functional API
Subclassing

19.2 Sequential API#

Use Sequential when layers form a simple linear stack.

🐍 Python
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(10,)),
    tf.keras.layers.Dense(32, activation="relu"),
    tf.keras.layers.Dense(1)
])

19.3 Functional API#

The Functional API allows more complex architectures.

🐍 Python
inputs = tf.keras.Input(shape=(10,))

x = tf.keras.layers.Dense(
    32,
    activation="relu"
)(inputs)

outputs = tf.keras.layers.Dense(1)(x)

model = tf.keras.Model(
    inputs=inputs,
    outputs=outputs
)

It is useful for:

  • Multiple inputs
  • Multiple outputs
  • Branching architectures
  • Skip connections
  • Shared layers

20. TensorFlow/Keras Training

20.1 compile()#

Before using Keras's standard training workflow, configure the model:

🐍 Python
model.compile(
    optimizer="adam",
    loss="mse",
    metrics=["mae"]
)

This specifies:

text
Optimizer
Loss
Metrics

20.2 fit()#

Training is commonly performed using:

🐍 Python
model.fit(
    X_train,
    y_train,
    epochs=10,
    batch_size=32,
    validation_data=(X_val, y_val)
)

The high-level flow is:

text
Training Data
     ↓
model.fit()
     ↓
Forward Pass
     ↓
Loss
     ↓
Gradients
     ↓
Optimizer
     ↓
Updated Weights

Keras manages much of this loop automatically.

20.3 evaluate()#

After training:

🐍 Python
model.evaluate(
    X_test,
    y_test
)

This computes the configured loss and metrics on the supplied data.

20.4 predict()#

To generate predictions:

🐍 Python
predictions = model.predict(X_test)

21. Keras Validation

Validation can be supplied directly to fit():

🐍 Python
history = model.fit(
    X_train,
    y_train,
    epochs=10,
    batch_size=32,
    validation_data=(X_val, y_val)
)

Keras evaluates the model on the validation data during training.

You can inspect the training history:

🐍 Python
print(history.history.keys())

Typical entries include:

›loss
val_loss

and any metrics specified during compile().

22. GPU Training in TensorFlow

22.1 GPU Availability#

TensorFlow can detect supported GPUs available to the environment.

🐍 Python
print(tf.config.list_physical_devices("GPU"))

If a compatible GPU setup is available, TensorFlow can place supported operations on it.

22.2 Explicit Device Placement#

TensorFlow provides device scopes:

🐍 Python
with tf.device("/GPU:0"):
    x = tf.random.normal((1000, 1000))

In many normal Keras workflows, explicit device placement is unnecessary because TensorFlow handles device placement automatically.

22.3 Keras GPU Training#

A normal Keras training call can use an available GPU without changing the basic API:

🐍 Python
model.compile(
    optimizer="adam",
    loss="mse"
)

model.fit(
    X_train,
    y_train,
    epochs=10
)

The actual GPU behavior depends on the TensorFlow installation, available hardware, supported operations, and environment configuration.

23. PyTorch vs TensorFlow/Keras

ConceptPyTorchTensorFlow/Keras
Tensortorch.Tensortf.Tensor
Model basenn.Moduletf.keras.Model
Layer examplenn.Linearkeras.layers.Dense
DatasetDatasetCommonly arrays/tensors or tf.data pipelines
DataLoaderDataLoadertf.data.Dataset
Automatic differentiationAutogradtf.GradientTape
TrainingOften custom loopmodel.fit() or custom loop
Optimizerstorch.optimtf.keras.optimizers
Lossestorch.nn lossestf.keras.losses
GPU.to(device)Automatic placement / device APIs
Model parametersstate_dict()Weights/model saving APIs

24. PyTorch Training Flow vs Keras Training Flow

PyTorch#

PyTorch commonly exposes the training process explicitly:

text
Dataset
   ↓
DataLoader
   ↓
Batch
   ↓
model()
   ↓
Loss
   ↓
loss.backward()
   ↓
optimizer.step()

You have direct control over the training loop.

Keras#

Keras provides a higher-level training interface:

text
Data
 ↓
model.fit()
 ↓
Forward Pass
 ↓
Loss
 ↓
Gradients
 ↓
Optimizer
 ↓
Updated Model

Much of the training loop is handled internally.

25. A Complete PyTorch Example

🐍 Python
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader


# Data
X = torch.randn(1000, 10)
y = torch.randn(1000, 1)

dataset = TensorDataset(X, y)

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True
)


# Model
class RegressionModel(nn.Module):

    def __init__(self):
        super().__init__()

        self.network = nn.Sequential(
            nn.Linear(10, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )

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


model = RegressionModel()


# Loss and optimizer
loss_fn = nn.MSELoss()

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


# Training
for epoch in range(10):

    model.train()

    total_loss = 0.0

    for X_batch, y_batch in loader:

        optimizer.zero_grad()

        predictions = model(X_batch)

        loss = loss_fn(
            predictions,
            y_batch
        )

        loss.backward()

        optimizer.step()

        total_loss += loss.item()

    print(
        f"Epoch {epoch + 1}, "
        f"Loss: {total_loss / len(loader):.4f}"
    )

The complete flow is:

text
TensorDataset
     ↓
DataLoader
     ↓
Batch
     ↓
Model
     ↓
Prediction
     ↓
MSE Loss
     ↓
Autograd
     ↓
Adam
     ↓
Updated Parameters

26. A Complete TensorFlow/Keras Example

🐍 Python
import tensorflow as tf


# Data
X = tf.random.normal((1000, 10))
y = tf.random.normal((1000, 1))


# Model
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(10,)),
    tf.keras.layers.Dense(
        64,
        activation="relu"
    ),
    tf.keras.layers.Dense(1)
])


# Configure training
model.compile(
    optimizer=tf.keras.optimizers.Adam(
        learning_rate=0.001
    ),
    loss=tf.keras.losses.MeanSquaredError(),
    metrics=[
        tf.keras.metrics.MeanAbsoluteError()
    ]
)


# Train
history = model.fit(
    X,
    y,
    epochs=10,
    batch_size=32,
    validation_split=0.2
)

The complete flow is:

text
TensorFlow Tensors
        ↓
Keras Model
        ↓
compile()
        ↓
fit()
        ↓
Forward Pass
        ↓
Loss
        ↓
Gradient Computation
        ↓
Adam
        ↓
Updated Weights

27. Important PyTorch Concepts to Remember

text
Tensor
  ↓
Dataset
  ↓
DataLoader
  ↓
nn.Module
  ↓
Forward Pass
  ↓
Loss Function
  ↓
Autograd
  ↓
Optimizer
  ↓
Parameter Update

For training:

🐍 Python
optimizer.zero_grad()
prediction = model(x)
loss = loss_fn(prediction, y)
loss.backward()
optimizer.step()

For validation:

🐍 Python
model.eval()

with torch.no_grad():
    prediction = model(x)

28. Important TensorFlow/Keras Concepts to Remember

text
TensorFlow Tensor
       ↓
Keras Layer
       ↓
Keras Model
       ↓
compile()
       ↓
fit()
       ↓
evaluate()
       ↓
predict()

The most important high-level APIs are:

🐍 Python
model.compile(...)
model.fit(...)
model.evaluate(...)
model.predict(...)

29. Summary Table

ConceptPyTorchTensorFlow/Keras
Tensor computationtorch.Tensortf.Tensor
Automatic differentiationAutogradGradient-based APIs such as GradientTape
Neural-network modelnn.Moduletf.keras.Model
Dataset abstractionDatasettf.data / other input sources
BatchingDataLoadertf.data.Dataset
Training loopCommonly explicitmodel.fit() or custom loop
Optimizertorch.optimtf.keras.optimizers
Losstorch.nn.*Losstf.keras.losses
GPU.to(device)Automatic device placement
Save model statestate_dict()Keras/TensorFlow saving APIs
High-level trainingMore manual by defaultmodel.fit()

30. Quick Recap

text
                    DEEP LEARNING FRAMEWORKS
                              │
               ┌──────────────┴──────────────┐
               │                             │
            PyTorch                    TensorFlow
               │                             │
      ┌────────┼────────┐                    │
      │        │        │                 Keras
   Tensor   Autograd  nn.Module               │
      │        │        │              ┌──────┼──────┐
      │        │        │              │      │      │
   Dataset  Training  Model          Layers Models Training
      │       Loop      │
DataLoader            Optimizer
                      Loss
                      GPU
                      Save/Load
                      Checkpoints

PyTorch mental model#

text
Tensor
  ↓
Dataset
  ↓
DataLoader
  ↓
nn.Module
  ↓
Prediction
  ↓
Loss
  ↓
Autograd
  ↓
Optimizer
  ↓
Updated Parameters

TensorFlow/Keras mental model#

text
TensorFlow Tensor
       ↓
Keras Layers
       ↓
Keras Model
       ↓
compile()
       ↓
fit()
       ↓
Loss + Gradients
       ↓
Optimizer
       ↓
Updated Weights

The core difference to remember is that PyTorch commonly gives you more direct control over the training loop, while Keras provides a high-level training interface that can manage much of the loop for you. Both frameworks provide the underlying building blocks needed for complete deep-learning workflows.

Knowledge Checkpoint

26. PyTorch Architecture Checkpoint

Q1.What is the difference between a PyTorch Tensor's `data` buffer and its `stride` tuple?
AThe data buffer is a contiguous 1D array of bytes in memory; stride specifies the number of memory elements to skip to move one step along each dimension.
BStride is the learning rate of the tensor.
CData buffer only exists on GPU while stride exists on CPU.
DStride determines the precision (FP16 vs FP32).
Q2.What is the purpose of `torch.no_grad()` context manager during inference evaluation?
AIt disables autograd tracking and avoids saving intermediate activation tensors to the computational graph, drastically reducing VRAM and execution time.
BIt forces all outputs to be positive.
CIt sets all weights to zero.
DIt compiles Python bytecode to C++.
Q3.How does `torch.utils.data.DataLoader` optimize mini-batch loading across CPU and GPU hardware?
AIt uses multiple background worker processes (`num_workers > 0`) to asynchronously fetch, transform, and pin CPU memory (`pin_memory=True`) for fast non-blocking CUDA transfers.
BIt deletes datasets after reading.
CIt converts image files to strings.
DIt trains model weights on the CPU before GPU transfer.
Track Your Learning

Finished studying this notebook?

Mark this guide as completed to update your course progress roadmap.

Previous Lesson
25. Parameter-Efficient Fine-Tuning (PEFT, LoRA & QLoRA)
Next Lesson
27. Model Training, Validation & Evaluation Pipelines