记录一次小白的tensorflow学习过程,也为有同样困扰的小白留下点经验。
先说我出错和解决的过程。在做风格迁移实验时,使用预加载权重的VGG19网络正向提取中间层结果,结果因为代码不当,在遍历图片提取时内存持续增长,导致提取几十个图片的特征内存就满了。
原因是在对每一张图片正向传播结束后,都会在留下中间信息。具体地说是在我将正向传播的代码与模型的代码分离了,在每次遍历图片时都会正向传播,在tensorflow中新增加了很多的计算节点(如tf.matmul等等),导致内存中遗留了大量的过期信息。
纠正的做法就是一个前提:避免在循环训练图片时额外使用tf计算资源。
使用placeholder作为输入数据的入口,在模型中定义需要使用的函数,包括正向传播。不要在遍历图片时额外使用tf计算。
遇到这种问题一定要回头检查代码,尤其是在别人写的代码基础上改时。 多学习公开的源码。
错误示例:
def
build_model(
model_path):
model_input
= tf.placeholder(
'float32', [
1,
IMAGE_HEIGHT,
IMAGE_WIDTH,
COLOR_CHANNELS])
vec1
=
...
......
return model_input,vec1
def
get_style_represent(
vec):
# 一些tf计算操作
return new_vec
with tf.Session()
as sess:
sess.run(tf.global_variables_initializer())
img_input,vec1
= build_model(
VGG19_MODEL)
# 加载模型
for cur_img_path
in imgs_path_list:
# 遍历图片
cur_img
= load_image(cur_img_path)
vec1_out
= sess.run(vec1,
feed_dict
= {img_input:cur_img})
# 正向传播输出模型中的vec1
# 对vec1进行一些处理,此处在遍历图片时额外使用了tensorflow的计算节点,导致在内存中遗留信息
new_vec
= get_style_represent(vec1_out)
正确示例:
def
build_model(
model_path):
model_input
= tf.placeholder(
'float32', [
1,
IMAGE_HEIGHT,
IMAGE_WIDTH,
COLOR_CHANNELS])
vec1
=
...
......
new_vec
=
...
# 将get_style_represent计算操作定义在模型中
return model_input,vec1,new_vec
with tf.Session()
as sess:
sess.run(tf.global_variables_initializer())
img_input,vec1,new_vec
= build_model(
VGG19_MODEL)
for cur_img_path
in imgs_path_list:
cur_img
= load_image(cur_img_path)
# 一次正向传播将处理后的vec1也得到了,避免在每次图片正向传播时留下额外信息
vec1_out,new_vec_out
= sess.run([vec1,new_vec],
feed_dict
= {img_input:cur_img})