Tensor shapes and dimensions in PyTorch - Model Pipeline Trace

Model Pipeline - Tensor shapes and dimensions

This pipeline shows how data moves through a simple PyTorch model, focusing on how tensor shapes change at each step. Understanding tensor shapes helps us avoid errors and design models correctly.

Data Flow - 5 Stages

1Input Data

64 rows x 3 channels x 32 height x 32 width→Raw image batch loaded→64 rows x 3 channels x 32 height x 32 width

Batch of 64 RGB images, each 32x32 pixels

↓

2Convolution Layer

64 x 3 x 32 x 32→Apply 16 filters with kernel size 3x3, stride 1, padding 1→64 x 16 x 32 x 32

Each image now has 16 feature maps of size 32x32

↓

3Max Pooling

64 x 16 x 32 x 32→2x2 max pooling with stride 2→64 x 16 x 16 x 16

Feature maps downsampled to 16x16

↓

4Flatten

64 x 16 x 16 x 16→Flatten each sample to a vector→64 x 4096

Each image represented as a 4096-length vector

↓

5Fully Connected Layer

64 x 4096→Linear layer to 10 output classes→64 x 10

Output logits for 10 classes per image

Training Trace - Epoch by Epoch

Epoch 1: 1.85 *****
Epoch 2: 1.20 ****
Epoch 3: 0.85  ***
Epoch 4: 0.65  **
Epoch 5: 0.50  *

Epoch	Loss ↓	Accuracy ↑	Observation
1	1.85	0.35	Loss starts high, accuracy low as model begins learning
2	1.20	0.55	Loss decreases, accuracy improves as model learns features
3	0.85	0.70	Model continues to improve, loss drops steadily
4	0.65	0.78	Good convergence, accuracy nearing 80%
5	0.50	0.85	Model well trained, loss low and accuracy high

Prediction Trace - 5 Layers

Layer 1: Input Tensor

Layer 2: Convolution Layer

Layer 3: Max Pooling

Layer 4: Flatten

Layer 5: Fully Connected Layer

Model Quiz - 3 Questions

What is the shape of the tensor after the convolution layer?

A 64 x 16 x 32 x 32

B 64 x 3 x 32 x 32

C 64 x 16 x 16 x 16

D 64 x 4096

Key Insight

Understanding tensor shapes at each step helps us design models that work without errors. Each layer changes the shape in a predictable way, and tracking these changes ensures data flows correctly through the model.