Batch size and shuffling in PyTorch - Model Pipeline Trace

Model Pipeline - Batch size and shuffling

This pipeline shows how data is split into batches and shuffled before training a simple model. Batch size controls how many samples the model sees at once, and shuffling mixes data to help the model learn better.

Data Flow - 3 Stages

1Raw dataset

1000 rows x 10 columns→Original dataset with 1000 samples and 10 features each→1000 rows x 10 columns

[[0.5, 1.2, ..., 0.3], [1.1, 0.7, ..., 0.9], ...]

↓

2Shuffle dataset

1000 rows x 10 columns→Randomly reorder all samples to mix data→1000 rows x 10 columns

[[1.1, 0.7, ..., 0.9], [0.5, 1.2, ..., 0.3], ...]

↓

3Batch creation

1000 rows x 10 columns→Split data into batches of 100 samples each→10 batches x 100 rows x 10 columns

[Batch 1: 100 samples, Batch 2: 100 samples, ...]

Training Trace - Epoch by Epoch

Loss
1.0 | *       
0.8 |  *      
0.6 |   *     
0.4 |    *    
0.2 |     *   
0.0 +---------
      1 2 3 4 5
      Epochs

Epoch	Loss ↓	Accuracy ↑	Observation
1	0.85	0.60	Initial training with shuffled batches, loss starts high, accuracy moderate
2	0.65	0.72	Loss decreases, accuracy improves as model learns from shuffled batches
3	0.50	0.80	Continued improvement, showing benefit of batch training and shuffling
4	0.40	0.85	Loss lowers further, accuracy rises, model converging well
5	0.35	0.88	Training stabilizes with good accuracy, showing effective batch size and shuffling

Prediction Trace - 4 Layers

Layer 1: Input batch

Layer 2: Model forward pass

Layer 3: Loss calculation

Layer 4: Backpropagation and update

Model Quiz - 3 Questions

Why do we shuffle the dataset before creating batches?

A To increase batch size automatically

B To mix data so the model sees varied samples in each batch

C To reduce the number of samples

D To make the model train faster by skipping samples

Key Insight

Shuffling data before batching helps the model learn from a good mix of samples each time. Choosing the right batch size balances training speed and stability, improving model accuracy over time.