安裝：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())