TensorFlow的几点小知识

最新推荐文章于 2022-06-21 00:57:58 发布

原创最新推荐文章于 2022-06-21 00:57:58 发布 · 2.8k 阅读

CC 4.0 BY-SA版权

文章标签：

25 篇文章

订阅专栏

9 篇文章

订阅专栏

本文介绍如何在TensorFlow中配置GPU资源，包括调节GPU内存占比和按需分配等技巧；此外还介绍了如何在训练过程中指定训练部分变量、实现权值共享以及指定GPU运行任务的方法。

1、调节GPU占比

TensorFlow比较贪心，默认会占用全部的GPU的资源。可以通过以下方式调节：

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.9
session = tf.Session(config=config)

另外，可以按需分配GPU资源

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)

2、寻找网络中的变量

在训练网络的过程中，可能只需要训练部分参数，固定其他参数，那怎么用TF实现呢？

我们可以在定义参数的时候采用

with tf.variable_scope('var'):

这一结构，相当于这些变量都是在scope 'var'下面。然后使用

theta = tf.get_collection(tf.GraphKeys.TRAINING_VARIABLES, scope='var')

就可以提取这些变量，在训练的时候，指定这些变量作为训练变量即可。

optimizer = ly.optimize_loss(loss=your_loss, learning_rate=your_learning_rate, optimizer=tf.train.AdamOptimizer, variables=theta, ...)

3、权值共享

我们可以使用

with tf.variable_scope('var'):

来实现这一功能。假设A,B,C权值共享，先定义A的网络结构和scope，然后定义B,C相同的网络结构，使用和A相同的scope，同时将reuse设置为True

with tf.variable_scope('var', reuse=True):

4、多GPU中，指定GPU运行TF

CUDA_VISIBLE_DEVICES=0 python -i train.py

未完待续。。。