在较少的内存中加载模型检查点

时间:2020-11-12 14:28:28

标签: python-3.x tensorflow memory-management pytorch out-of-memory

我有一个问题,我找不到在线答案。我已经训练了一个模型,其检查点文件大约为 20Gb 。由于我的系统没有足够的RAM(或者Colab / Kaggle-限制为16Gb),因此我无法将模型用于预测。

我知道必须将模型加载到内存中才能进行推理,但是有人知道一种解决方法或方法可以:

  1. 保存一些内存并能够将其加载到16Gb RAM(用于CPU)中,或者将其加载到TPU / GPU中
  2. 可以使用任何框架(因为我将同时使用两者)Tensorflow + Keras或PyTorch(我现在正在使用)

在这两个库中甚至都可以使用这种方法吗?我的顽强解决方案之一可能不是分块加载它,实质上是为模型权重和偏差保留了一个缓冲区,并据此执行了计算-尽管我还没有找到实现的方法。

我还想补充一点,我不介意性能下降,因为低规格的硬件可能会出现这种情况。只要不超过2周:)我绝对可以等那么长时间....

1 个答案:

答案 0 :(得分:1)

您可以尝试以下操作:

  1. 将模型分为两部分
  2. 将权重加载到两个部分,分别调用Composable
  3. 使用您的输入调用第一个模型
  4. 使用第一个模型的输出调用第二个模型