Slow numpy array indexing for keras time series generator

Viewed 639

I use the keras time series generator for training a neural network with LSTM cells, which unfortunately proved to be a bottleneck in training.

Below is a simplified example to run, which shows the high runtime of the batch generator. It is important to note that the rows from the dataset are chosen randomly and thus a sliding window is not possible. During the training the CPUs are running continuously at about 80%, whereas the GPU is running at a single-digit percentage rate.

def get_time_series(data, index, look_back, batch_size):
    samples1 = np.empty((batch_size, look_back, np.size(data, axis=1)))
    rows = np.random.randint(look_back, np.size(data, axis=1), size=batch_size)
    for j, row in enumerate(rows):
        indices = range(rows[j] - look_back, rows[j], 1)
        samples1[j] = data[indices]
    return samples1


data = np.random.rand(100000, 20)
start = time.time()
batch = get_time_series(data, index=50, look_back=1000, batch_size=2**12)
print("Batch generator needs",  time.time()-start,  "seconds")

Result:

Batch generator needs 0.6224319934844971 seconds

I already tried to build the 3-d array first, so I only have to index the array-rows in the *get_time_series-*Function. This was about 60 times faster during the training, but leads to an "out of memory error" with large datasets.

Does anyone have ideas on how to improve the performance of this bottleneck? Work with pointer, faster indexing methods, ...

Thanks, Max

1 Answers

EDIT 2:

Not sure if this is going to be any faster, but you can also just do something like this. It still relies on advanced indexing, although over contiguous data, so maybe it's a bit better?:

import numpy as np

def get_time_series(data, indices, look_back):
    # Make sure indices are big enough
    indices = indices[indices >= look_back]
    # Make indexing matrix
    idx = indices[:, np.newaxis] + np.arange(-look_back, 0)
    # Make batch
    return data[idx]

You would use it for example like this:

import numpy as np

def get_time_series(data, indices, look_back):
    indices = indices[indices >= look_back]
    idx = indices[:, np.newaxis] + np.arange(-look_back, 0)
    return data[idx]

def make_batches(data, look_back, batch_size):
    indices = np.random.permutation(np.arange(look_back, len(data) + 1))
    for i in range(0, len(indices), batch_size):
        yield get_time_series(data, indices[i:i + batch_size], look_back)

data = ...
look_back = ...
batch_size = ...
for batch in make_batches(data, look_back, batch_size):
    # Use batch

EDIT:

If you want to shuffle the examples, you could first make the sliding window for the whole dataset (which should not take any memory or time) and then take batches from a shuffled index:

# Make sliding window with the previous function
data_sw = get_time_series(data, 0, look_back, len(data))
# Random index
batch_idx = np.random.permutation(len(data_sw))
# To get the first batch
batch = data_sw[batch_idx[:batch_size]]

I think this does what you want, and should be quite faster than using loops:

import numpy as np

def get_time_series(data, index, look_back, batch_size):
    from numpy.lib.stride_tricks import as_strided
    # Index should be at least as big as look_back to have enough elements before it
    index = max(index, look_back)
    # Batch size should not go beyond the array
    batch_size = min(batch_size, len(data) - index + 1)
    # Relevant slice for the batch
    data_slice = data[index - look_back:index + batch_size]
    # Reshape with stride tricks as a "sliding window"
    data_strides = data_slice.strides
    batch_shape = (batch_size, look_back, data_slice.shape[-1])
    batch_strides = (data_strides[0], data_strides[0], data_strides[1])
    return as_strided(data_slice, batch_shape, batch_strides, writeable=False)

# Test
data = np.arange(300).reshape((100, 3))
batch = get_time_series(data, 20, 5, 4)
print(batch)

Output:

[[[45 46 47]
  [48 49 50]
  [51 52 53]
  [54 55 56]
  [57 58 59]]

 [[48 49 50]
  [51 52 53]
  [54 55 56]
  [57 58 59]
  [60 61 62]]

 [[51 52 53]
  [54 55 56]
  [57 58 59]
  [60 61 62]
  [63 64 65]]

 [[54 55 56]
  [57 58 59]
  [60 61 62]
  [63 64 65]
  [66 67 68]]]
Related