Re: [问题] tensorflow- 训练在gpu 预测在cpu

楼主: snakei14702 (sun)   2022-09-25 03:11:37
※ 引述《snakei14702 (阿中是我们的神)》之铭言:
: 最近在玩强化训练的时候遇到一点小问题就是
: 因为模型的深度很深需要GPU来跑
: 但是在run env 计算(中间有用到model.predict)一些东西的时候
: 如果用GPU跑就显得太慢没效率...
: 所以想要请教一下各位有没有办法在fit model时用GPU
: 然后用model.prict的时候用cpu呢?
: 简化的代码如下:
: for i in range(1000):
: model.fit(x,y)
: model.predict(x)
: 基本上我有用过
: for i in range(1000):
: os.environ["CUDA_VISIBLE_DEVICES"]="0"
: model.fit(x,y)
: os.environ["CUDA_VISIBLE_DEVICES"]="-1"
: model.predict(x)
: and
: for i in range(1000):
: with tf.device('/GPU:0'):
: model.fit(x,y)
: with tf.device('/CPU:0'):
: model.predict(x)
: 这两个方法都没效= =
: 有点难过
: 不知道版上强者没有研究过,可以指导一下小弟
: 非常感谢!!
谢谢各位强者的帮忙
后来研究一下
用@tf.fuction修饰
然后用model(x)取代model.predict(x)
就可以跑的爆炸快了!!!
以前觉得@tf.fuction好像也没多厉害
不过在RL玩游戏的时候整个让我很惊艳!
再次感谢各位的帮忙!
作者: ToastBen (吐司边)   2022-09-26 09:29:00

Links booklink

Contact Us: admin [ a t ] ucptt.com