使用 CPU 加载模型权重时内存不足

Question

1.I 使用

在 GPU 上训练了我的深度学习模型

import keras    
model.save('trained_model.h5')

2. 我想加载模型并使用

   预测CPU

   模型 = load_model('trained_model.h5')

但是出现以下错误

tensorflow.python.framework.errors_impl.InternalError: CUDA runtime implicit initialization on GPU:0 failed. Status: out of memory

如何在 CPU 上加载经过训练的模型？

Answer 1

您得到的错误代码是说您的 GPU 内存不足。要在 CPU 上 运行，您需要输入

CUDA_VISIBLE_DEVICES=""

在终端或代码中将是

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'