torch Conv1d小输入反向传播时GPU内存异常激增问题求助

问题现象 #

使用torch.Conv1d时遇到非常怪异的现象：当输入样本数低于某一阈值时，反向传播阶段GPU内存占用会急剧飙升，幅度可达一个数量级甚至更高。推测这与PyTorch/CUDA根据输入维度和可用内存选择不同卷积算法有关，但该问题会在运行时引发棘手的OOM（内存不足）错误。

可复现代码 #

import torch
from torch import nn
%load_ext pytorch_memlab

base_ch = 512
d_pos = 64


def print_gpu_mem_usage(prefix=""):
    print(f"{prefix}Peak memory: {torch.cuda.max_memory_allocated() / 1024 ** 3:.2f} GB"
          f" | {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB"
          f" (Current: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB"
          f" | {torch.cuda.memory_reserved() / 1024 ** 3:.2f} GB)")


def isolated_conv(v):
    samp_conv = nn.Conv1d(base_ch + d_pos, 2 * base_ch, kernel_size=1, padding='valid').cuda()
    mn = samp_conv(v).mean()
    mn.backward()

测试结果 #

样本数5000时的内存占用

执行命令：

%mlrun -f isolated_conv isolated_conv(torch.rand(5000, base_ch+d_pos, 11).cuda())

内存监控结果：

样本数4000时的内存占用

执行命令：

%mlrun -f isolated_conv isolated_conv(torch.rand(4000, base_ch+d_pos, 11).cuda())

内存监控结果：

补充测试信息 #

• 调换两种输入的测试顺序，问题依然存在。
• 测试运行于Docker环境，可提供Dockerfile用于复现。
• 尝试设置torch.backends.cudnn.deterministic=True，仅略微改变阈值，当样本数降至3000时仍会出现内存暴增问题。

环境信息 #

• torch == 2.0.1
• pytorch-memlab == 0.3.0
• Nvidia 2080Ti
• 驱动版本：525.105.17
• CUDA版本：12.0

问题原因分析 #

该现象源于CuDNN的卷积算法自动选择机制：CuDNN会根据输入张量尺寸、硬件性能、可用内存等因素，在直接卷积、Winograd算法、FFT卷积等多种实现中自动选择最优方案。对于小批量输入，CuDNN可能会选择计算速度快但内存开销极大的算法（如部分Winograd变体或未优化的直接卷积反向实现），导致反向传播时内存占用暴增。

解决方案 #

1. 强制指定CuDNN卷积算法
   关闭CuDNN的自动基准测试，限制算法选择范围，避免内存开销过大的实现：

   # 初始化模型前添加
   torch.backends.cudnn.benchmark = False
   torch.backends.cudnn.deterministic = True
   
   # 或针对特定卷积层单独设置
   conv = nn.Conv1d(base_ch + d_pos, 2 * base_ch, kernel_size=1, padding='valid').cuda()
   with torch.backends.cudnn.flags(enabled=True, deterministic=True, benchmark=False):
       mn = conv(v).mean()
       mn.backward()
   

2. 限制CuDNN工作空间
   设置CuDNN可使用的最大工作空间，迫使它选择内存友好型算法：

   # 设置最大工作空间为1GB（单位：字节）
   torch.backends.cudnn.set_max_workspace_size(1 * 1024**3)
   

3. 升级PyTorch版本
   PyTorch 2.0.x的CuDNN集成可能存在算法选择bug，升级到2.1或更高版本，官方大概率已修复小批量卷积反向传播的内存异常问题。

4. 替换1x1卷积实现
   对于kernel_size=1的卷积，完全可以用矩阵乘法替代，绕过CuDNN的算法选择逻辑：

   def conv1x1(input, weight, bias=None):
       # input shape: [B, C_in, L]
       # weight shape: [C_out, C_in, 1]
       weight = weight.squeeze(-1)  # 转为[C_out, C_in]
       output = torch.matmul(input.transpose(1,2), weight.transpose(0,1)).transpose(1,2)
       if bias is not None:
           output += bias.unsqueeze(0).unsqueeze(-1)
       return output
   
   # 使用示例
   def isolated_conv(v):
       weight = torch.randn(2*base_ch, base_ch+d_pos, 1).cuda()
       bias = torch.randn(2*base_ch).cuda()
       mn = conv1x1(v, weight, bias).mean()
       mn.backward()
   

内容的提问来源于Stack Exchange，提问作者Roi