【记录】keras强制使用cpu运行

原创于 2020-09-24 16:46:13 发布 · 1.4k 阅读

3 ·

CC 4.0 BY-SA版权

本文介绍如何在Keras代码中遇到显存不足时，通过设置CUDA_VISIBLE_DEVICES为-1来让模型在CPU上运行。同时提供限制GPU显存使用的方法和配置TensorFlow会话的实例。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

问题描述：我有一台GPU机器，但是显存被占满了，剩余显存不足以支撑我的服务；所以我想让keras代码使用cpu运行；

方法很简单，在程序当中添加如下代码，这个命令本来是用来指定GPU设备的，GPU设备编号从0开始，编码-1表示使用CPU；

import os
 
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

CPU使用率限制：num_cores表示使用的cpu核数；

num_cores = 2
config = tf.ConfigProto(intra_op_parallelism_threads=num_cores,
                        inter_op_parallelism_threads=num_cores,
                        allow_soft_placement=True, )
session = tf.Session(config=config)
K.set_session(session)

扩展：

如果要限制GPU显存使用：

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = .4
set_session(tf.Session(config=config))

🍅

💀