Alperen AYDIN
Alperen AYDIN

Reputation: 567

Move graph trained the GPU to be tested on the CPU

So I have this CNN which I train on the GPU. During the training, I regularly save checkpoint.

Later on, I want to have a small script that reads .meta file and the checkpoint and do some tests on a CPU. I use the following the code:

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
with sess.as_default():
    with tf.device('/cpu:0'):
        saver = tf.train.import_meta_graph('{}.meta'.format(model))
        saver.restore(sess,model)

I keep getting this error which tell me that the saver is trying to put the operation on the GPU.

How can i change that?

Upvotes: 0

Views: 728

Answers (2)

Surajit Podder
Surajit Podder

Reputation: 73

Move all the ops to CPU using _set_device API. https://github.com/tensorflow/tensorflow/blob/r1.14/tensorflow/python/framework/ops.py#L2255

with tf.Session() as sess: 
    g = tf.get_default_graph() 
    ops = g.get_operations() 
    for op in ops: 
        op._set_device('/device:CPU:*')

Upvotes: 1

Yaroslav Bulatov
Yaroslav Bulatov

Reputation: 57923

Hacky work-around, open your graph definition file (ending with .pbtxt), and remove all lines starting with device:

For programmatic approach you can see how TensorFlow exporter does this with clear_devices although that uses regular Saver, not meta graph exporter

Upvotes: 0

Related Questions