个人笔记-TensorFlow中GPU及模型加载

一.GPU运行Tensorflow

1.在运行之前先查看GPU的使用情况

nvidia-smi
watch nvidia-smi 实时查看GPU

2.查看是否可以用GPU

# 查看GPU
from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

3.指定GPU运行

训练代码中加入 os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘0’
备注：使用 GPU 0，这里的‘0’是根据查看可用GPU里面的数字来的。

4.GPU占用量设置

1.设置定量的GPU显存使用量:

config = tf.ConfigProto()

config.gpu_options.per_process_gpu_memory_fraction = 0.8 # 占用GPU 80%的显存

session = tf.Session(config=config)

2.动态申请显存（在服务器跑建议）

gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

二.模型加载

1.保存和加载模型

推荐看这个: 『TensorFlow』模型保存和载入方法汇总.

2.加载部分权重

# 不加载w_1,b_1
ckpt = tf.train.latest_checkpoint(models_path)
exclude = ['w_1','b_1',]
variables_to_restore = slim.get_variables_to_restore(exclude=exclude)
saver = tf.train.Saver(variables_to_restore)
sess.run(tf.global_variables_initializer()) 
saver.restore(sess, ckpt)

3.测试遇到的问题

加载模型继续训练后，保存下来的模型你会发现大小比你开始加载那个模型要小一点，在测试的时候可能会发生 NotFoundError: Key xxxx not found in checkpoint 这个错误。

#在加载模型成功后训练时，保存模型前需要加一句 saver = tf.train.Saver()
saver = tf.train.Saver(max_to_keep=None)  
saver.save(sess, save_path=models_path, global_step=i)             

第一篇博客，如有问题，欢迎指正 。

如有疑问，加群97359887交流