Keras以及Tensorflow强制使用GPU

请叫我算术嘉

于 2019-12-06 15:23:55 发布

阅读量8.5k

点赞数 1

CC 4.0 BY-SA版权

分类专栏： Tensorflow 文章标签： tensorflow keras gpu

本文链接：https://blog.youkuaiyun.com/ssjdoudou/article/details/103423769

Tensorflow 专栏收录该内容

7 篇文章

订阅专栏

本文探讨了在TensorFlow 2.0环境中配置GPU的方法，详细介绍了使用Keras进行设置时遇到的问题及解决方案，同时提供了一个可行的替代方案，即通过设置CUDA_VISIBLE_DEVICES环境变量来指定使用的GPU。

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

环境：python3.6+tensorflow==2.0+keras2.3.1

方法一：

import tensorflow as tf
import keras.backend.tensorflow_backend as KTF
 
KTF.set_session(tf.Session(config=tf.ConfigProto(device_count={'gpu':0})))

这里在tensorflow2.0中肯定报错
修改如下：

import tensorflow as tf
import keras.backend.tensorflow_backend as KTF
 
KTF.set_session(tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(device_count={'gpu':0})))

然而这样还是不行，就是告诉你tensorflow2.0中不能这样用
遂放弃

方法二

CUDA_VISIBLE_DEVICES=0 python3 xxx.py

貌似可行