python-3.x - 以更少的内存加载模型检查点

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

我有一个问题,在网上找不到任何答案。我训练了一个模型,其检查点文件约为 20 GB。由于我的系统没有足够的 RAM(或 Colaboratory/Kaggle - 限制为 16 GB),因此我无法使用我的模型进行预测。

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

  1. 节省一些内存并能够将其加载到 16 GB RAM(用于 CPU)或 TPU 中的内存中/GPU
  2. 可以使用任何框架(因为我会同时使用这两个框架)TensorFlow + Keras ,或PyTorch (我现在正在使用)

这样的方法是否可以在这两个库中实现?我的尝试性解决方案之一可能不是分块加载它,本质上是为模型权重和偏差维护一个缓冲区并相应地执行计算 - 尽管我还没有找到任何实现。

I would also like to add that I wouldn't mind the performance slowdown since it is to be expected with low-specification hardware. As long as it doesn't take more than two weeks :) I can definitely wait that long...

最佳答案

你可以尝试以下方法:

  1. 将模型分为两部分
  2. 分别调用 model.load_weights(by_name=True) 将权重加载到两个部分
  3. 使用您的输入调用第一个模型
  4. 使用第一个模型的输出调用第二个模型

关于python-3.x - 以更少的内存加载模型检查点,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/64805828/

相关文章:

c++ - struct rebind::other 是什么意思?

c - 通过/proc/$PID/smaps 在 Linux 内核中进行脏页统计

Python写入文件

python-3.x - 如何修复 subprocess.check_output() 中的挂起

python - 如何避免使用 Python 多处理在 fork 进程中加载​​父模块

python - TensorFlow:有没有办法测量模型的 FLOPS?

c - 如何释放临时字符串的内存?

python-3.x - 基于语义相似度的聚类不返回值

python - Pandas 线图(按月),按带有时间戳的 SQL 导出的行业分组

python - 为什么损失函数中使用.numpy()时Tensorflow的自动微分失败?