🎓 A1Learner
🔍

PyTorch Vectorization over Loops

Vectorization is doing maths on a whole tensor at once instead of looping row by row. One batched call hands all the rows to fast compiled code in a single step.

The loop, timed

The loop sums one row at a time and asks Python for each value. Ten thousand small steps add up.

import time

import torch

torch.manual_seed(0)
data = torch.randn(10000, 16)

start = time.perf_counter()
total = 0.0
for i in range(data.shape[0]):
    total += data[i].sum().item()
print(f"Loop over 10000 rows: {(time.perf_counter() - start) * 1000:.0f} ms")

One call does every row

A single sum call walks the whole tensor inside compiled code. The total matches the loop's total.

start = time.perf_counter()
total_fast = data.sum().item()
print(f"One sum over all rows: {(time.perf_counter() - start) * 1000:.0f} ms")
print("Same total:", abs(total - total_fast) < 0.01)

Example

import time

import torch

torch.manual_seed(0)
data = torch.randn(10000, 16)

start = time.perf_counter()
loop_rows = []
for i in range(data.shape[0]):
    loop_rows.append((data[i] * 2 + 1).mean())
loop_result = torch.stack(loop_rows)
loop_ms = (time.perf_counter() - start) * 1000

start = time.perf_counter()
batch_result = (data * 2 + 1).mean(dim=1)
batch_ms = (time.perf_counter() - start) * 1000

print("Rows scored:", loop_result.shape[0])
print("Results identical:", torch.allclose(loop_result, batch_result))
print(f"Python loop: {loop_ms:.0f} ms")
print(f"One batched op: {batch_ms:.0f} ms")

Key points

  • Each loop step pays Python's own step cost before the maths starts
  • Use one operation on the whole tensor to replace a loop over rows
  • A batched operation runs the maths in compiled code across all rows
  • allclose checks that two results agree within a small tolerance
  • The speed ratio changes from machine to machine
  • In short, vectorization trades many small Python steps for one compiled pass
  • A loop is still the right tool when each row needs different work

In this example, a table of 10,000 rows is scored two ways. The loop visits each row, scales it, and averages it. The batched version runs the same maths as one mean over each row. An allclose check prints True, so the speedup costs no correctness. The loop took about 25 times longer on the recorded run. That ratio varies by machine, but the gap stays large because the loop pays Python's step cost 10,000 times.

Note

Vectorization wins when every row gets the same maths.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What makes the batched version faster on the same maths?