pytorch计算模型的显存占用率和节省内存技巧

版权声明:本文为博主原创文章,未经博主允许不得转载。 https://blog.csdn.net/hajungong007/article/details/86558193

计算模型的显存占用率

节省内存的技巧

  • 减小Batch-size;
  • 减小图片尺寸;
  • 将计算过程分为多步,保存中间结果再计算后面一半的模型,参考此文
  • 使用软batch训练model,即通过修改优化算法,是网络进行N次前向传播,但只进行一次反向传播;
  • 使用pooling,减小特征图的size;
  • 减少全连接层的使用。relu(inplace=true),inplace_abn使用半精度float16;
  • optimizer的变换使用,理论上,sgd<momentum<adam,可以从计算公式中看出有额外的中间变量;
  • Depthwise Convolution;
  • 尽可能使用inplace操作, 比如relu 可以使用 inplace=True;
  • 不要一次性把数据加载进来,使用 pytorch 的 Dataset api;
    大致程序如下:

    from torch.utils.data import DataLoader
    data_train_loader = DataLoader(data_train, batch_size=64, shuffle=True)
    for batch_idx, samples inenumerate(data_train_loader):
    net(samples)

优化pytorch代码

猜你喜欢

转载自blog.csdn.net/hajungong007/article/details/86558193