一句话实现自动调整batch_size，再也不会Cuda out of memory

防搞活机

已于 2023-01-11 12:29:08 修改

阅读量1.8k

点赞数 3

于 2023-01-09 20:42:37 首次发布

本文链接：https://blog.youkuaiyun.com/qq_51750957/article/details/128620616

版权

文章介绍了如何利用huggingface的accelerate库中的find_executable_batch_size装饰器，自动调整模型验证循环的batch_size，以防止CUDA内存溢出。通过装饰函数，可以避免手动调整并确保程序在显存不足时仍能运行。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

需要huggingface的accelerate库

核心函数（装饰器）
accelerate.find_executable_batch_size(function: callable = None, starting_batch_size: int = 128):
官方文档
（很短，小白基本看不懂，实际只需改2句话）

传统代码

以一个torch的valid_loop函数为例，valid_loop函数需要以batch_size为第一个参数（必须位置上是第一个，不能用关键字传参）

# 验证loop，验证一个epoch
def valid_loop_fn(batch_size: int, dataset, model, data_collator=None):
    model.eval()
    losses = []
    data_loader = DataLoader(dataset, batch_size=batch_size,
                             shuffle=False, collate_fn=data_collator)

    with torch.no_grad():
        tk0 = tqdm(data_loader, total=len(data_loader), desc='valid')
        for step, batch in enumerate(tk0):
            outputs = model(**batch)
            loss = outputs[0]
            losses.append(loss.cpu().item())

    # 计算平均loss
    avg_loss = np.mean(losses)
    return avg_loss

正常情况下，如果传入一个较大的batch_size例如512，就会报错

valid_losses = valid_loop_fn(512,tokenized_dataset['test'], model, data_collator)

然后RuntimeError: CUDA out of memory
然后就要手动修改batch_size，很麻烦

改进代码

此时，只需要accelerate.find_executable_batch_size对valid_loop进行装饰即可，即函数前面加一句@accelerate.find_executable_batch_size(starting_batch_size=512)
如果显存不够，会自动将batch_size减半，不会报错

import accelerate
# 验证loop，验证一个epoch
@accelerate.find_executable_batch_size(starting_batch_size=128)
def valid_loop_fn(batch_size: int, dataset, model, data_collator=None):
    model.eval()
    losses = []
    data_loader = DataLoader(dataset, batch_size=batch_size,
                             shuffle=False, collate_fn=data_collator)

    with torch.no_grad():
        tk0 = tqdm(data_loader, total=len(data_loader), desc='valid')
        for step, batch in enumerate(tk0):
            outputs = model(**batch)
            loss = outputs[0]
            losses.append(loss.cpu().item())

    # 计算平均loss
    avg_loss = np.mean(losses)
    return avg_loss

另一个改动，就是调用valid_loop时不需要再传入batch_size，如果传入反而会报错。因为装饰器已经帮你设定了batch_size，你只需要告诉被装饰的函数除了batch_size以外的信息即可，如下

valid_losses = valid_loop_fn(tokenized_dataset['test'], model, data_collator)

经过一点微小的改动，再也不会因为batch_size太大而报错了
(*^▽^*)