当模型已经足够小时减少 GPU 内存使用

数据挖掘 张量流 显卡 硬件
2021-10-12 06:52:47

我训练了一个模型,冻结成一个PB(protocol buffer)文件和一些变量的目录,总大小大约31M。我们使用 GPU 卡部署它并遵循这个答案并将其设置per_process_gpu_memory_fraction为一个非常小的数字以使内存约为 40M。该程序执行得非常好,但是当我们检查 GPU 使用情况时nvidia-smi,显示内存使用量约为 500M。

那么我的问题是如何证明这个差距是合理的?我们怎样才能减少它?我们可以做一些类似量化的事情来减少 500M 吗?我们想将它部署到边缘设备中,所以 500M 太大了。

1个回答

关于差距,我无法回答。我可以帮助您的是如何减少内存使用量。

您可以尝试的一件事是您可以将权重转换为 fp16 并使用 fp16 进行推理。不要担心你不会失去显着的数值精度,特别是因为你正在使用它进行部署。事实上,这就是人们在混合精度训练中进行前向传递的方式。从理论上讲,这可以提高速度并减少内存使用量。希望能帮助到你。