如何根据随机生成的批量大小从数据集中切分不同长度的切片？

需求说明 #

需要从数据集中提取随机大小的样本，具体流程：

1. 定义各样本大小的出现概率，生成一批随机的批次大小（如batch_sizes = [3,4,4,5,2]）
2. 按照这批批次大小，依次从数据集中切分对应长度的子样本，示例效果：

   [0,1,2], [3,4,5,6], [7,8,9,10], ...
   

问题痛点 #

之前尝试用固定步长的循环切分，代码如下：

for i in range(0, len(data)-batch_size+1, batch_size):
    batch = data[i:i+batch_size]
    print('Batch: ', batch)

但该方法仅适用于固定批次大小，无法适配动态变化的batch_sizes列表。

解决方案 #

通过维护一个当前起始索引，遍历每个批次大小，动态调整切分位置即可实现需求，代码示例：

from random import choices

# 1. 定义样本大小及权重，生成随机批次列表
population = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
weights = [0.025, 0.025, 0.05, 0.1, 0.2, 0.3, 0.1, 0.1, 0.05, 0.05]

batch_sizes = []
for _ in range(5):
    n = int(choices(population, weights)[0])
    batch_sizes.append(n)

# 2. 模拟待切分的数据集
data = list(range(1000004))  # 对应0到1000003的数据集

# 3. 按动态批次大小切分数据
current_idx = 0
batches = []
for size in batch_sizes:
    # 处理剩余数据不足当前批次的情况
    end_idx = current_idx + size
    if end_idx > len(data):
        batch = data[current_idx:]
    else:
        batch = data[current_idx:end_idx]
    batches.append(batch)
    current_idx = end_idx

# 输出结果
for idx, batch in enumerate(batches):
    print(f"批次 {idx+1}: {batch}")

核心逻辑 #

• 用current_idx记录每次切分的起始位置，初始值为0
• 遍历每个批次大小size，计算结束位置current_idx + size
• 若结束位置超出数据集长度，直接取剩余所有数据；否则取对应切片
• 更新current_idx为当前结束位置，进入下一次循环

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