GPU显存不足-优化方案-pytorch

按照优化成本，对可以减少显存占用方案进行排序：

1.计算总loss时候进行代码优化

loss本身是一个包含梯度信息的 tensor，正确的求损失和的方式为：

total_loss += loss.item()

2.降低batch_size

适当降低batch size， 则模型每层的输入输出就会成线性减少， 效果相当明显。必须要多个batch_size时候可以采用梯度累计的方式，但是这比上直接降低batch_size占用显存会大点。

传统方式：

for i,(feature,target) in enumerate(train_loader):
    outputs = model(feature)  # 前向传播
    loss = criterion(outputs,target)  # 计算损失

    optimizer.zero_grad()   # 清空梯度
    loss.backward()  # 计算梯度
    optimizer.step()  # 反向传播， 更新网络参数

梯度累计：

for i,(features,target) in enumerate(train_loader):
    outputs = model(images)  # 前向传播
    loss = criterion(outputs,target)  # 计算损失
    loss = loss/accumulation_steps   # 可选，如果损失要在训练样本上取平均

    loss.backward()  # 计算梯度
    if((i+1)%accumulation_steps)==0:
        optimizer.step()        # 反向传播，更新网络参数
        optimizer.zero_grad()   # 清空梯度

3.Relu 的 inplace 参数

激活函数 Relu() 有一个默认参数 inplace ，默认为Flase， 当设置为True的时候，我们在通过relu() 计算得到的新值不会占用新的空间而是直接覆盖原来的值，这表示设为True， 可以节省一部分显存。

4.释放不需要的张量和变量

采用del释放你不再需要的张量和变量，这也要求我们在写模型的时候注意变量的使用，不要随心所欲，漫天飞舞。

5.数据变小

图像适当的缩小可以极大的降低参数量。

6.精简模型

减少卷积核数量； 尽量少的使用 Linear 等。全连接层参数较多，较少参数或则不用全连接层。使用全局平均池化进行替代。

参考链接.