Tensorflow, restore variables in a specific device

前端 未结 2 2004
萌比男神i
萌比男神i 2021-01-14 03:33

Maybe my question is a bit naive, but I really didn\'t find anything in the tensorflow documentation.

I have a trained tensorflow model where the variables of it was

相关标签:
2条回答
  • 2021-01-14 04:00

    Use clear_devices flag, ie

    saver = tf.train.import_meta_graph("/tmp/graph.meta", clear_devices=True)
    
    0 讨论(0)
  • 2021-01-14 04:09

    I'm using tensorflow 0.12 and clear_devices=True and tf.device('/cpu:0') was not working with me (saver.restore was still trying to assign variables to /gpu:0).

    I really needed to force everything to /cpu:0 since I was loading several models which wouldn't fit in GPU memory anyways. Here are two alternatives to force everything to /cpu:0

    1. Set os.environ['CUDA_VISIBLE_DEVICES']=''
    2. Use the device_count of ConfigProto like tf.Session(config=tf.ConfigProto(device_count={"GPU": 0, "CPU": 1}))
    0 讨论(0)
提交回复
热议问题