keras以tensorflow指定GPU与CPU

最新推荐文章于 2022-11-08 21:25:22 发布

WhatUwannadO

最新推荐文章于 2022-11-08 21:25:22 发布

阅读量2k

点赞数

CC 4.0 BY-SA版权

分类专栏： Deep Learning

本文链接：https://blog.youkuaiyun.com/WhatUwannadO/article/details/84560492

Deep Learning 专栏收录该内容

5 篇文章

订阅专栏

本文介绍了如何在TensorFlow和Keras中指定使用GPU或CPU进行计算，通过代码示例展示了如何在资源受限的情况下合理分配GPU和CPU资源，以及如何在Keras中构建和训练一个简单的神经网络。

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

由于显卡显存不够,所以就只能一部分模型用gpu,一部分用cpu.

# tensorflow指定cpu或gpu
# import tensorflow as tf
# import os

# from tensorflow.python.client import device_lib
# print(device_lib.list_local_devices())
#
# print(tf.__version__)
# print(tf.__path__)
#
# # 新建一个 graph.
## with tf.device("/gpu:0"):
# with tf.device("/cpu:0"):
#
#
#     a = tf.constant(5.0)
#     b = tf.constant(6.0)
#     c = a * b
#     # 新建session with log_device_placement并设置为True.
#     sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True,
#                                         log_device_placement=True))
#     # 运行这个 op.
#     print (sess.run(c))
#     #任务完成 关闭会话
#     sess.close()


# keras指定cpu或gpu
from keras.models import Sequential
from keras.layers import Dense
import numpy as np

import tensorflow as tf
import keras.backend.tensorflow_backend as KTF

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
# The GPU id to use, usually either "0" or "1"
# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# 使用CPU
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
# 观察使用的设备是cpu还是gpu
KTF.set_session(tf.Session(config=tf.ConfigProto(log_device_placement=True)))
# 输入训练数据 keras接收numpy数组类型的数据
x = np.array([[0, 1, 0],
              [0, 0, 1],
              [1, 3, 2],
              [3, 2, 1]])
y = np.array([0, 0, 1, 1]).T
# 最简单的序贯模型，序贯模型是多个网络层的线性堆叠
simple_model = Sequential()
# dense层为全连接层
# 第一层隐含层为全连接层 5个神经元 输入数据的维度为3
simple_model.add(Dense(5, input_dim=3, activation='relu'))
# 第二个隐含层 4个神经元
simple_model.add(Dense(4, activation='relu'))
# 输出层为1个神经元
simple_model.add(Dense(1, activation='sigmoid'))
# 编译模型,训练模型之前需要编译模型
# 编译模型的三个参数：优化器、损失函数、指标列表
simple_model.compile(optimizer='sgd', loss='mean_squared_error')
# 训练网络 2000次
# Keras以Numpy数组作为输入数据和标签的数据类型。训练模型一般使用fit函数
simple_model.fit(x, y, epochs=2000)
# 应用模型 进行预测
y_ = simple_model.predict_classes(x[0:1])
print("[0,1,0]的分类结果：" + str(y[0]))