[TensorFlow] How to generate the Memory Report from Grappler?

In the previous post, I introduce the way to generate cost and model report from Grappler.
https://danny270degree.blogspot.com/2019/01/tensorflow-how-to-generate-cost-and.html
In this post, I will continue to introduce the memory report which I think that is very useful. Please refer to my previous post to get the model code.



First, we can define a virtual GPU device's properties as follows:

from tensorflow.core.protobuf import device_properties_pb2
from tensorflow.python.grappler import item

# You can define your GPU device's properties
grappler_item = item.Item(m)
device_properties = device_properties_pb2.DeviceProperties(
    type='GPU',
    frequency=1582,
    num_cores=3584,
    environment={
   'architecture': '8'
})
named_device = device_properties_pb2.NamedDevice(
    properties=device_properties, name='/device:GPU:0')

grappler_cluster = gcluster.Cluster(
    disable_detailed_stats=False,
    disable_timeline=False,
    devices=[named_device])

Generate Memory Report:

peak_usage = grappler_cluster.DeterminePeakMemoryUsage(grappler_item)
detailed_report = True

for device, snapshot in peak_usage.items():
    report = ""
    peak_usage = snapshot[0]
    report += "Peak usage for device " + device + ": " + str(
        peak_usage) + " bytes\n"
    if detailed_report:
        live_tensors = snapshot[1]
        for tensor in live_tensors:
            op_name = tensor[0]
            output_id = tensor[1]
            mem_used = tensor[2]
            report += "  " + str(op_name) + ":" + str(output_id) + " uses " + str(mem_used) + " bytes\n"
    print(report)

Result: ( you can see the peak memory usage in your model with the batch size)

Peak usage for device /device:GPU:0: 725364 bytes
  gradients/MatMul_grad/MatMul_1:0 uses 31360 bytes
  Variable/read:0 uses 31360 bytes
  gradients/add_grad/tuple/control_dependency_1:0 uses 40 bytes
  gradients/add_grad/tuple/control_dependency:0 uses 4000 bytes
  Variable_1:0 uses 40 bytes
  Variable:0 uses 31360 bytes
  Placeholder:0 uses 313600 bytes
  GradientDescent/learning_rate:0 uses 4 bytes
  gradients/MatMul_grad/MatMul:0 uses 313600 bytes

So, what if I change my batch size in model from 100 to 1000?

Ans: the peak memory usage is from 725364 bytes to 6454864 bytes, which increases 8.9 times.

Peak usage for device /device:GPU:0: 6454864 bytes

  gradients/add_grad/Reshape:0 uses 40000 bytes

  gradients/add_grad/Reshape_1:0 uses 40 bytes

  gradients/add_grad/tuple/control_dependency:0 uses 40000 bytes

  gradients/add_grad/tuple/control_dependency_1:0 uses 40 bytes

  gradients/MatMul_grad/MatMul:0 uses 3136000 bytes

  gradients/add_grad/Sum:0 uses 4 bytes

  gradients/sparse_softmax_cross_entropy_loss/xentropy/xentropy_grad/mul:0 uses 40000 bytes

  gradients/add_grad/BroadcastGradientArgs:1 uses 0 bytes

  gradients/add_grad/Sum_1:0 uses 4 bytes

  Variable/read:0 uses 31360 bytes

  gradients/add_grad/Shape_1:0 uses 4 bytes

  gradients/add_grad/Shape:0 uses 8 bytes

  Variable_1:0 uses 40 bytes

  Variable:0 uses 31360 bytes

  Placeholder:0 uses 3136000 bytes

  GradientDescent/learning_rate:0 uses 4 bytes