How to break numpy array into smaller chunks/batches, then iterate through them
consider array a
a = np.array([[1, 2, 3],
[4, 5, 6],
[7, 8, 9],
[10, 11, 12]])
Option 1
use reshape
and //
a.reshape(a.shape[0] // 2, -1, a.shape[1])
array([[[ 1, 2, 3],
[ 4, 5, 6]],
[[ 7, 8, 9],
[10, 11, 12]]])
Option 2
if you wanted groups of two rather than two groups
a.reshape(-1, 2, a.shape[1])
array([[[ 1, 2, 3],
[ 4, 5, 6]],
[[ 7, 8, 9],
[10, 11, 12]]])
Option 3
Use a generator
def get_every_n(a, n=2):
for i in range(a.shape[0] // n):
yield a[n*i:n*(i+1)]
for sa in get_every_n(a, n=2):
print sa
[[1 2 3]
[4 5 6]]
[[ 7 8 9]
[10 11 12]]
You can use numpy.split
to split along the first axis n
times, where n
is the number of desired batches. Thus, the implementation would look like this -
np.split(arr,n,axis=0) # n is number of batches
Since, the default value for axis
is 0
itself, so we can skip setting it. So, we would simply have -
np.split(arr,n)
Sample runs -
In [132]: arr # Input array of shape (10,3)
Out[132]:
array([[170, 52, 204],
[114, 235, 191],
[ 63, 145, 171],
[ 16, 97, 173],
[197, 36, 246],
[218, 75, 68],
[223, 198, 84],
[206, 211, 151],
[187, 132, 18],
[121, 212, 140]])
In [133]: np.split(arr,2) # Split into 2 batches
Out[133]:
[array([[170, 52, 204],
[114, 235, 191],
[ 63, 145, 171],
[ 16, 97, 173],
[197, 36, 246]]), array([[218, 75, 68],
[223, 198, 84],
[206, 211, 151],
[187, 132, 18],
[121, 212, 140]])]
In [134]: np.split(arr,5) # Split into 5 batches
Out[134]:
[array([[170, 52, 204],
[114, 235, 191]]), array([[ 63, 145, 171],
[ 16, 97, 173]]), array([[197, 36, 246],
[218, 75, 68]]), array([[223, 198, 84],
[206, 211, 151]]), array([[187, 132, 18],
[121, 212, 140]])]