tensorflow 图大小

Question

我想知道是否有一种简单的方法可以在运行 tensorflow session 之前检查 tensorflow 图所需的大小/内存。

我正在寻找一些东西，我可以不断更改定义图形的系统参数，并可以看到图形相应地变得有多大(在内存中)。

Answer 1

我做了类似的事情，我想查看模型中的参数数量。

vars = 0
for v in tf.all_variables():
    vars += np.prod(v.get_shape().as_list())
print(vars)

现在 vars 包含图中所有变量维度的乘积之和。如果每个变量的类型为 tf.float32，则可以将 vars 乘以 4 以获得所有变量消耗的字节数。然而，这只是一个下限，并且会有一些额外的开销。另外，我认为计算梯度需要大量内存，因为它需要存储模型中每个点的激活以进行向后传递。