为什么tensorflow仅在单个图像上进行正向传递就分配这么多的gpu内存? 使用MobileNet和gpu_options.allow_growth = True重现此问题的简单代码:
# Load model
graph = tf.Graph() with graph.as_default():
image = tf.placeholder(tf.float32, shape=[ 224, 224, 3], name="input_image")
# Mobilenet preprocessing
image = tf.to_float(image)
image = tf.subtract(image, 0.5)
image = tf.multiply(image, 2.0)
batched_image = tf.expand_dims(image, 0)
with slim.arg_scope(mobilenet_v1.mobilenet_v1_arg_scope()):
logits, endpoints = mobilenet_v1.mobilenet_v1(batched_image, num_classes=1001,
is_training=False,
depth_multiplier=1.0)
predictions = tf.to_int32(tf.argmax(logits, 1))
init_fn = slim.assign_from_checkpoint_fn("mobilenet_v1_1.0_224.ckpt",
slim.get_variables_to_restore())
#Load image & preprocessing
img = cv2.imread('cat.jpg')
img = cv2.resize(img, (224, 224))
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(graph=graph, config=config) as sess:
#Initialize weights
init_fn(sess)
for i in range(10000):
c = sess.run(predictions, feed_dict={"input_image:0":img})
网络仅在一张图像上进行预测(批量大小= 1),但tensorflow仍会分配7800 MB的gpu内存。即使将MobileNet深度乘数设置为0.25,tensorflow仍会分配3GB。
我试图冻结模型权重并剥离所有训练操作,但是内存消耗仍然很高。
将模型与TF服务一起使用会减少内存消耗吗? tensorRT呢?
预先感谢