在这种情况下我只使用CPU:
当我编写emb函数时,我在tf.device中添加了(' / cpu:0')
任何人都可以告诉我何时应该使用CPU,何时应该使用GPU?
def emb(inputs, vocabulary_size, embedding_size,
W_init=tf.random_uniform_initializer(-0.1, 0.1),
W_regu=None,
reuse=None,
scope=None,
printable=True):
with tf.variable_scope(scope, default_name=func_name(), reuse=reuse):
with tf.device('/cpu:0'):
emb_mat = slim.model_variable('emb_mat',
shape=[vocabulary_size, embedding_size],
initializer=W_init,
regularizer=W_regu)
outputs = tf.nn.embedding_lookup(emb_mat, inputs)
if printable and not tf.get_variable_scope().reuse: print(
"[info] utils.emb: {0} {1}->{2}".format(outputs.name, inputs.get_shape(), outputs.get_shape()))
return outputs