具有批量大小的 Numpy 切片

Question

我有一个 numpyA的 shape数组(550,10)。我的批量大小为 100，即我想要从中获取多少数据行A。在每次迭代中，我想从 A 中提取 100 行。但是当我到达最后 50 行时，我想要从 A 中提取最后 50 行和前 50 行。

我有一个这样的函数：

def train(index, batch_size):

    if(batch_size + index < A.shape(0)):
          data_end_index = index + batch_size
          batch_data = A[index:batch_end_index,:]
    else:
          data_end_index = index + batch_size - A.shape(0) #550+100-600 = 50
          batch_data = A[500 to 549 and 0 to 49] # How to slice here ?

如何执行最后一步？

Answer 1

你可以试试：

import numpy as np
data=np.random.rand(550,10)
batch_size=100

for index in range(0,data.shape[0],batch_size):
    batch=data[index:min(index+batch_size,data.shape[0]),:]
    print(batch.shape)

输出：

(100, 10)
(100, 10)
(100, 10)
(100, 10)
(100, 10)
(50, 10)