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

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

目录

• • 1.计算总loss时候进行代码优化
  • 2.降低batch_size
  • 3.Relu 的 inplace 参数
  • 4.释放不需要的张量和变量
  • 5.数据变小
  • 6.精简模型

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 等。全连接层参数较多，较少参数或则不用全连接层。使用全局平均池化进行替代。

参考链接.

本文来自互联网用户投稿，文章观点仅代表作者本人，不代表本站立场，不承担相关法律责任。如若转载，请注明出处。 如若内容造成侵权/违法违规/事实不符，请点击【内容举报】进行投诉反馈！