python - 如何将GPU上训练的模型加载到CPU(系统)内存中？

Question

我在 GPU 上训练了一个模型，现在我尝试在 CPU 上评估它(GPU 用于不同的训练运行)。但是，当我尝试使用以下方式加载它时:

with tf.device('/cpu:0'):
    model = tf.keras.models.load_model('models/resnet50_untrained.h5', compile=False)

我收到一个CUDA_ERROR_OUT_OF_MEMORY:

2019-10-17 09:25:23.088408: W tensorflow/compiler/xla/service/platform_util.cc:256] unable to create StreamExecutor for CUDA:0: failed initializing StreamExecutor for CUDA device ordinal 0: Internal: failed call to cuDevicePrimaryCtxRetain: CUDA_ERROR_OUT_OF_MEMORY: out of memory; total memory reported: 7981694976
2019-10-17 09:25:23.088505: F tensorflow/stream_executor/lib/statusor.cc:34] Attempting to fetch value instead of handling error Internal: no supported devices found for platform CUDA

(我也尝试设置 compile=True 得到相同的结果。)

该模型似乎正在加载到已被另一个实例使用的 GPU 中。如何强制 keras/tensorflow 将其加载到系统内存中并在 CPU 上执行？

Answer 1

您可以使用 tf.device('/cpu:0'): 定义内部的所有内容吗？除了库导入部分和测试之外。

如果这不起作用，那么创建一个虚拟环境并安装普通的tensorflow而不是gpu版本，然后尝试它。如果仍然是 OOM 错误，则说明 CPU 正在使用，没有足够的内存来加载这个训练好的模型。