如何按固定步长遍历4D矩阵第一维度并拆分批次获取数据

原代码核心错误 #

• 索引逻辑错误：Python序列索引从0开始，a[1:20:5]的切片会直接丢弃索引为0的第一条数据，且遍历得到的i是矩阵的实际元素值，不是索引，后续用a[i,:]索引会完全不符合预期
• 没有实现批次切分：每次循环只取了单条数据，无法得到每个批次5条的结构

正确实现代码 #

全程不调用reshape，每次直接切分第一维度的连续5条作为批次：

import numpy as np

output = []
# 按需求构造形状为(20,3,100,100)的输入矩阵
a = np.random.randint(0, 100, size=(20, 3, 100, 100))
# 按批次步长5遍历第一维度的起始位置
for batch_idx in range(0, a.shape[0], 5):
    # 切取当前批次的5条数据，形状为(5,3,100,100)可直接输入神经网络
    batch_data = a[batch_idx: batch_idx + 5, :, :, :]
    output.append(batch_data)
# 转换为numpy数组后即得到目标形状(4,5,3,100,100)
output = np.array(output)

验证方法 #

打印output.shape即可确认输出维度符合要求。

内容的提问来源于stack exchange，提问作者Saman