Выберите несколько ломтиков из массива Numpy одновременно

Question

Я хочу реализовать векторизованный алгоритм SGD и хотел бы генерировать несколько мини-пакетов одновременно.

Предположим, data = np.arange(0, 100), miniBatchSize=10, n_miniBatches=10 и indices = np.random.randint(0, n_miniBatches, 5) (5 мини-пакетов). Чего я хотел бы достичь, так это

miniBatches = np.zeros(5, miniBatchSize)
for i in range(5):
     miniBatches[i] = data[indices[i]: indices[i] + miniBatchSize]

Есть ли способ избежать цикла?

Спасибо!

Answer 1

Это можно сделать, используя трюки с шагами :

from numpy.lib.stride_tricks import as_strided

a = as_strided(data[:n_miniBatches], shape=(miniBatchSize, n_miniBatches), strides=2*data.strides, writeable=False)    
miniBatches = a[:, indices].T


# E.g. indices = array([0, 7, 1, 0, 0])
Output:
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [ 7,  8,  9, 10, 11, 12, 13, 14, 15, 16],
       [ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10],
       [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9]])