PyTorch自定义自适应分类层后出现GPU显存泄露问题求助

问题根因 #

你每次调用classifierAdaptive函数时，都会重新实例化一个新的nn.Linear层并移动到CUDA设备，PyTorch默认会将所有新创建的层参数加入计算图，且不会主动释放，每轮前向传播都会新增一份线性层参数占用显存，迭代几次就会直接占满显存。同时你每次创建的线性层参数没有被复用，也没有加入模型的参数列表管理，除了显存溢出问题外，参数也无法正常训练更新。

修复方案 #

方案1：用自适应池化固定输出维度（最推荐，真正适配任意输入尺寸） #

不需要动态创建线性层，在特征提取后加一层自适应平均/最大池化，把输出的空间维度固定为你需要的6*6，直接复用原来的静态分类器即可：
首先在__init__中定义自适应池化层和原分类器：

self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
self.classifier = nn.Sequential(
    nn.Linear(512 * 6 * 6, 1024),
    nn.LeakyReLU(0.2, True),
    nn.Linear(1024, 1),
    nn.Sigmoid()
)

然后修改forward逻辑：

def forward(self, x):
    out = self.features(x)
    out = self.avgpool(out) # 不管输入特征的空间维度是多少，输出统一为[B, 512, 6, 6]
    out = torch.flatten(out, 1)
    out = self.classifier(out)
    return out

该方案不需要动态创建任何层，显存占用和原有代码一致，同时完美支持任意尺寸的输入patch。

方案2：缓存动态创建的Linear层（如果确实需要动态匹配输入维度） #

如果你的输入patch尺寸只有固定几种，不是完全随机，可以把创建过的线性层存在模型的成员变量里复用，避免每次新建：

def __init__(self):
    super().__init__()
    self.linear_cache = nn.ModuleDict() # 用ModuleDict管理动态创建的层，会被模型正常识别参数
    self.lrelu = nn.LeakyReLU(0.2, True)
    self.dens2 = nn.Linear(1024, 1)
    self.sig = nn.Sigmoid()

def classifierAdaptive(self,img):
    out = torch.flatten(img, 1)
    in_dim = 512 * img.shape[2] * img.shape[3]
    # 只在第一次遇到对应维度的时候创建层，后续复用
    if str(in_dim) not in self.linear_cache:
        self.linear_cache[str(in_dim)] = nn.Linear(in_dim, 1024).cuda()
    out = self.linear_cache[str(in_dim)](out)
    out = self.lrelu(out)
    out = self.dens2(out)
    out = self.sig(out)
    return out

注意如果输入维度种类过多，该方案还是会额外占用显存，通用性不如方案1。

额外优化点 #

• 你之前写的out = outLinear(out).cuda()是冗余操作，outLinear已经在CUDA上，输出张量本身就在CUDA设备上，不需要再调用cuda()，多余操作会增加显存拷贝开销。
• 动态创建的层必须用nn.ModuleList/nn.ModuleDict存储，否则模型无法正常识别参数进行反向传播更新，也无法正常保存加载权重。

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