python - Tensorflow:在 pb 模型中使用 tensorflow.contrib.memory_stats.MaxBytesInUse

Question

我正在尝试获取 tensorflow 模型的内存使用情况，从卡住的 pb 文件加载模型:

import tensorflow as tf

def load_graph_def(model_filepath):
    # Expects frozen graph in .pb format
    with tf.gfile.GFile(model_filepath, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    return graph_def


def load_graph(graph_def):
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name="")
    return graph

def print_memory_usage(sess):
    def convert_units(size_bytes):
        import math
        if size_bytes == 0:
            return "0B"
        size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
        i = int(math.floor(math.log(size_bytes, 1024)))
        p = math.pow(1024, i)
        s = round(size_bytes / p, 2)
        return "%s %s" % (s, size_name[i])

    from tensorflow.contrib.memory_stats import MaxBytesInUse
    print("Max memory usage: ", convert_units(sess.run(MaxBytesInUse())))


model_filepath = './inception_v3_2016_08_28_frozen.pb'
graph_def = load_graph_def(model_filepath)
graph = load_graph(graph_def)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(graph=graph, config=config)
print_memory_usage(sess)

这会产生错误:

WARNING: Logging before flag parsing goes to stderr.
W1109 23:25:09.024481 140071828465472 deprecation_wrapper.py:119] From pb_memory_usage.py:5: The name tf.gfile.GFile is deprecated. Please use tf.io.gfile.GFile instead.

W1109 23:25:09.024652 140071828465472 deprecation_wrapper.py:119] From pb_memory_usage.py:6: The name tf.GraphDef is deprecated. Please use tf.compat.v1.GraphDef instead.

W1109 23:25:09.454025 140071828465472 deprecation_wrapper.py:119] From pb_memory_usage.py:34: The name tf.ConfigProto is deprecated. Please use tf.compat.v1.ConfigProto instead.

W1109 23:25:09.454198 140071828465472 deprecation_wrapper.py:119] From pb_memory_usage.py:36: The name tf.Session is deprecated. Please use tf.compat.v1.Session instead.

2019-11-09 23:25:09.454362: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-11-09 23:25:09.459323: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcuda.so.1
2019-11-09 23:25:09.659849: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x46edbb0 executing computations on platform CUDA. Devices:
2019-11-09 23:25:09.659873: I tensorflow/compiler/xla/service/service.cc:175]   StreamExecutor device (0): GeForce GTX 1080 Ti, Compute Capability 6.1
2019-11-09 23:25:09.659880: I tensorflow/compiler/xla/service/service.cc:175]   StreamExecutor device (1): GeForce GTX 1080 Ti, Compute Capability 6.1
2019-11-09 23:25:09.661890: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 3468535000 Hz
2019-11-09 23:25:09.662429: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x52564b0 executing computations on platform Host. Devices:
2019-11-09 23:25:09.662447: I tensorflow/compiler/xla/service/service.cc:175]   StreamExecutor device (0): <undefined>, <undefined>
2019-11-09 23:25:09.664178: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1640] Found device 0 with properties:
name: GeForce GTX 1080 Ti major: 6 minor: 1 memoryClockRate(GHz): 1.582
pciBusID: 0000:01:00.0
2019-11-09 23:25:09.664992: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1640] Found device 1 with properties:
name: GeForce GTX 1080 Ti major: 6 minor: 1 memoryClockRate(GHz): 1.582
pciBusID: 0000:02:00.0
2019-11-09 23:25:09.665185: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcudart.so.10.0
2019-11-09 23:25:09.666079: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcublas.so.10.0
2019-11-09 23:25:09.666819: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcufft.so.10.0
2019-11-09 23:25:09.667017: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcurand.so.10.0
2019-11-09 23:25:09.668051: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcusolver.so.10.0
2019-11-09 23:25:09.668813: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcusparse.so.10.0
2019-11-09 23:25:09.671132: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcudnn.so.7
2019-11-09 23:25:09.675589: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1763] Adding visible gpu devices: 0, 1
2019-11-09 23:25:09.675617: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcudart.so.10.0
2019-11-09 23:25:09.677460: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1181] Device interconnect StreamExecutor with strength 1 edge matrix:
2019-11-09 23:25:09.677473: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1187]      0 1
2019-11-09 23:25:09.677478: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1200] 0:   N Y
2019-11-09 23:25:09.677482: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1200] 1:   Y N
2019-11-09 23:25:09.680015: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1326] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10457 MB memory) -> physical GPU (device: 0, name: GeForce GTX 1080 Ti, pci bus id: 0000:01:00.0, compute capability: 6.1)
2019-11-09 23:25:09.681237: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1326] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:1 with 10479 MB memory) -> physical GPU (device: 1, name: GeForce GTX 1080 Ti, pci bus id: 0000:02:00.0, compute capability: 6.1)
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 303, in __init__
    fetch, allow_tensor=True, allow_operation=True))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 3796, in as_graph_element
    return self._as_graph_element_locked(obj, allow_tensor, allow_operation)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 3875, in _as_graph_element_locked
    raise ValueError("Tensor %s is not an element of this graph." % obj)
ValueError: Tensor Tensor("MaxBytesInUse:0", shape=(), dtype=int64) is not an element of this graph.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "pb_memory_usage.py", line 37, in <module>
    print_memory_usage(sess)
  File "pb_memory_usage.py", line 28, in print_memory_usage
    print("Max memory usage: ", convert_units(sess.run(MaxBytesInUse())))
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 950, in run
    run_metadata_ptr)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 1158, in _run
    self._graph, fetches, feed_dict_tensor, feed_handles=feed_handles)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 474, in __init__
    self._fetch_mapper = _FetchMapper.for_fetch(fetches)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 274, in for_fetch
    return _ElementFetchMapper(fetches, contraction_fn)
  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py", line 310, in __init__
    'Tensor. (%s)' % (fetch, str(e)))
ValueError: Fetch argument <tf.Tensor 'MaxBytesInUse:0' shape=() dtype=int64> cannot be interpreted as a Tensor. (Tensor Tensor("MaxBytesInUse:0", shape=(), dtype=int64) is not an element of this graph.)

Answer 1

您的 session 使用在 load_graph 函数中创建的图形。但是，MaxBytesInUse 运算符是在默认图上创建的(您可以使用 tf.get_default_graph() 获取)。这意味着运算符不是 session 使用的图形的一部分。

为了解决这个问题，您可以在创建运算符时指定正确的图形:

    from tensorflow.contrib.memory_stats import MaxBytesInUse
    with sess.graph.as_default():
        print("Max memory usage: ", convert_units(sess.run(MaxBytesInUse())))