如何在PyTorch中正确重置GPU内存以适配批量大小搜索？

我编写了一个通过二分法查找指定GPU所能支持的模型最大批量大小的函数，日志示例如下：

Batch size 1 succeeded. Increasing to 2...
Batch size 2 succeeded. Increasing to 4...
Batch size 4 succeeded. Increasing to 8...
Batch size 8 succeeded. Increasing to 16...
Batch size 16 succeeded. Increasing to 32...
Batch size 32 succeeded. Increasing to 64...
Batch size 64 succeeded. Increasing to 128...
Batch size 128 succeeded. Increasing to 256...
Batch size 256 succeeded. Increasing to 512...
Batch size 512 failed. Binary searching...
# We start with the bounds as (256 - 50) to 512 to detect the bug
# detailed later in this post
Batch size 359 failed. New bounds: [206, 359]
Batch size 282 failed. New bounds: [206, 282]
Batch size 244 failed. New bounds: [206, 244]
Batch size 225 failed. New bounds: [206, 225]
Batch size 215 failed. New bounds: [206, 215]
Batch size 210 failed. New bounds: [206, 210]
Batch size 208 failed. New bounds: [206, 208]
Batch size 207 failed. New bounds: [206, 207]

但日志出现异常：初始时批量大小256运行成功，但后续二分搜索阶段，更小的批量大小却运行失败。我认为这是GPU内存未被正确回收导致的问题。目前我在每次执行前向/反向传播函数前都会调用torch.cuda.empty_cache()，但这似乎不够。请问还需要执行哪些操作来重置GPU内存状态？

参考代码如下：

def binary_search_batch_size(cfg: Settings):
    def _is_cuda_oom(e: RuntimeError):
        """Determines if error is CUDA Out of Memory and if adaptive_grad_accum is enabled."""
        return 'CUDA out of memory' in str(e)

    cfg = Settings.parse_obj(cfg)
    model = ModelWrapper(model=cfg.model.model(), out_dim=cfg.model.output_dim).cuda()
    optimizer = cfg.optimizer.create_optimizer(model)

    def run():
        with torch.cuda.amp.autocast(enabled=True):
            torch.cuda.empty_cache()
            run_batch(model, optimizer, shape)

    batch_size = 1
    max_batch_size = 1
    while True:
        try:
            batch_size = max_batch_size
            shape = (batch_size, 3, 224, 224)
            run()
            max_batch_size *= 2
            print(f"Batch size {batch_size} succeeded. Increasing to {max_batch_size}...")
        except RuntimeError as e:
            if not _is_cuda_oom(e):
                raise e
            print(f"Batch size {batch_size} failed. Binary searching...")
            # the 50 acts as a bullshit check to make sure we haven't regressed somehow
            low = batch_size // 2 - 50
            high = batch_size
            while low + 1 < high:
                batch_size = (low + high) // 2
                shape = (batch_size, 3, 224, 224)
                try:
                    run()
                    low = batch_size
                    print(f"Batch size {batch_size} succeeded. New bounds: [{low}, {high}]")
                except RuntimeError as e:
                    if not _is_cuda_oom(e):
                        raise e
                    high = batch_size
                    print(f"Batch size {batch_size} failed. New bounds: [{low}, {high}]")
            max_batch_size = low
            break
    return max_batch_size

def run_batch(model, optimizer, shape):
    batch = dict(
        aug1=torch.randint(0, 256, shape, dtype=torch.uint8).cuda(),
        aug2=torch.randint(0, 256, shape, dtype=torch.uint8).cuda(),
    )
    
    image_logits1, image_logits2 = model(batch)
    loss_val = loss(image_logits1, image_logits2, 0)
    optimizer.zero_grad()
    loss_val.backward()
    optimizer.step()

解决方法 #

1. 显式释放张量并触发Python垃圾回收

torch.cuda.empty_cache()仅释放PyTorch已放弃引用但CUDA runtime仍持有的内存，若Python侧仍有张量引用，内存不会被回收。需在每次测试后主动删除大张量，并触发垃圾回收：

修改run_batch和run函数：

def run_batch(model, optimizer, shape):
    batch = dict(
        aug1=torch.randint(0, 256, shape, dtype=torch.uint8).cuda(),
        aug2=torch.randint(0, 256, shape, dtype=torch.uint8).cuda(),
    )
    
    image_logits1, image_logits2 = model(batch)
    loss_val = loss(image_logits1, image_logits2, 0)
    optimizer.zero_grad(set_to_none=True)  # 用set_to_none更彻底释放梯度内存
    loss_val.backward()
    optimizer.step()
    
    # 释放当前批次的所有张量
    del batch, image_logits1, image_logits2, loss_val
    torch.cuda.empty_cache()

def run():
    with torch.cuda.amp.autocast(enabled=True):
        run_batch(model, optimizer, shape)
    import gc
    gc.collect()  # 触发Python垃圾回收，释放未引用的对象
    torch.cuda.empty_cache()

2. 避免优化器状态累积

每次测试后，优化器的动量、权重缓存等会占用GPU内存。如果内存压力大，可以在每次二分搜索迭代前重置优化器（会增加初始化开销）：

# 在二分循环内每次测试前重新初始化优化器
optimizer = cfg.optimizer.create_optimizer(model)

3. 启用同步CUDA操作定位泄漏

设置环境变量CUDA_LAUNCH_BLOCKING=1，让CUDA操作同步执行，避免异步操作导致的内存泄漏难以追踪：

export CUDA_LAUNCH_BLOCKING=1

4. 修正二分搜索初始下界

日志中初始下界设为256-50=206不合理，应直接使用之前成功的256作为下界，减少不必要的内存测试：

low = batch_size // 2  # 替换原来的 batch_size //2 -50

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