如何控制AWS SageMaker批量转换的批量大小？解决内存溢出问题

一、批次大小混杂的原因 #

• SageMaker Batch Transform的分片逻辑结合MaxPayloadInMB与SplitType（默认Line）实现：
  1. 先按设置的MaxPayloadInMB大小拆分文件内容，若拆分点恰好落在某行文本中间，会将整行归到前一个批次，导致该批次行数略超出理论值；
  2. 若输入包含多个文件，小文件会被优先合并至接近MaxPayloadInMB的批次，剩余未合并的小文件会单独成批；
  3. 你的测试集25k样本共2.5MB，设MaxPayloadInMB为1MB时理论上拆分为3个批次，但实际出现大小混杂，大概率是输入文件的存储结构（比如包含多个小文件）或行分隔边界处理逻辑，导致部分批次被合并/拆分，最终形成不同大小的批次。

二、突破最小批量限制的方案 #

由于MaxPayloadInMB最小仅能设为1MB，无法直接控制到更小的批次，可通过自定义推理脚本（inference.py）实现二次批量拆分：

1. 在predict函数中，接收SageMaker传来的大批次数据后，手动拆分为自定义大小的子批次（比如每个子批次1000条）；
2. 逐个处理子批次，避免一次性加载过多数据触发内存溢出；
3. 合并所有子批次的推理结果后返回给SageMaker。

示例代码片段：

def predict(data, model):
    # 假设data为包含所有样本的列表/数组
    sub_batch_size = 1000  # 自定义子批次大小
    results = []
    # 循环拆分并处理子批次
    for i in range(0, len(data), sub_batch_size):
        sub_batch = data[i:i+sub_batch_size]
        sub_result = model(sub_batch)
        results.extend(sub_result)
    return results

三、额外优化建议 #

• 微调MaxConcurrentTransforms：若内存压力仍大，可适当减少worker数量，避免多个worker同时加载大批次导致内存耗尽；
• 规范输入文件格式：确保所有输入文件为行分隔的纯文本，避免非标准格式干扰分片逻辑；
• 添加内存监控：在推理脚本中引入psutil库记录内存使用情况，精准调整子批次大小。

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