Why is stacking a list of small tensors faster than updating a slice of a large tensor?

Viewed 142

I want to update slices of a tensor sequentially, basically iterating over multiple time steps. Naively, I thought that reserving the memory upfront and then updating a slice in that memory was faster than concatenating/stacking a list thereof in the end, but it seems as if I'm wrong. Any ideas why that is?

import torch

batch_size = 256
time_steps = 50
channels = 3
width = 20
heigth = 20
device = 'cuda'

def reserve_memory():
    out_state = torch.zeros((batch_size, time_steps, channels, width, heigth), device=device)
    state = torch.ones((batch_size, channels, width, heigth), device=device)
    for step in range(time_steps):
        state = state * 1e-3
        out_state[:,step] = state

def stack_memory():
    out_state = []
    state = torch.ones((batch_size, channels, width, heigth), device=device)
    for step in range(time_steps):
        state = state * 1e-3
        out_state.append(state)
    out_state = torch.stack(out_state, 1)
%timeit -r 100 reserve_memory()

1.04 ms ± 49 µs per loop (mean ± std. dev. of 100 runs, 1 loop each)

%timeit -r 100 stack_memory()

707 µs ± 9.09 µs per loop (mean ± std. dev. of 100 runs, 1000 loops each)

0 Answers
Related