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)