官网地址:https://pypi.org/project/memory_profiler/
安装:pip install -U memory_profiler
import tensorflow as tf
from memory_profiler import profile
a = tf.constant([[1,2,3]])
b = tf.Variable([[4], [5], [6]])
d = tf.matmul(a, b)
@profile()
def main():
with tf.Session() as sess:
for i in range(1000):
sess.run(tf.global_variables_initializer())
@profile()
def adds(a, b):
print(a+b)
if __name__=="__main__":
adds(a,b)
main()
输出:
Line # Mem usage Increment Line Contents
================================================
116 147.5 MiB 147.5 MiB @profile()
117 def adds(a, b):
118 147.5 MiB 0.0 MiB return a+b
Line # Mem usage Increment Line Contents
================================================
109 147.5 MiB 147.5 MiB @profile()
110 def main():
111 148.2 MiB 0.7 MiB with tf.Session() as sess:
112 189.8 MiB 0.2 MiB for i in range(1000):
113 189.8 MiB 40.7 MiB sess.run(tf.global_variables_initializer())