最近在使用Tensorflow搭建神经网络,某些时候网络复杂时在本地运行代码会出现显存不足错误。想要继续运行代码有两种方法,一是将代码上传至云端如kaggle,colab上运行。二是切换为CPU运行。

        由于有关Tensorflow和GPU使用问题已经反复出现过很多次,每次都需要重新搜索。在这里记录一下,方便以后查阅。

使用CPU运行代码

        只需要在需要运行代码前加上以下代码。

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

使用GPU运行代码

        有时在切换为CPU后希望重新使用GPU,需要添加以下代码。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

        此外,在切换GPU/CPU运算前,需要检查电脑的硬件情况。如果没有GPU就不存在使用GPU计算的情况了。查看GPU可以运行以下代码。

gpu = tf.test.gpu_device_name()
print(gpu)

        若输出为空,则代表电脑不具有可以进行Tensorflow运算的GPU。

Logo

基于 Vue 的企业级 UI 组件库和中后台系统解决方案,为数万开发者服务。

更多推荐