如何分配推理所需的最小GPU内存?

时间:2018-12-03 16:22:03

标签: python tensorflow

为什么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呢?

预先感谢

0 个答案:

没有答案