utensor_cgen_mem_alloc.py

from utensor_cgen.backend.graph_lower import TensorAllocationPlanner
from utensor_cgen.frontend.tensorflow import GraphDefParser

parser = GraphDefParser(config={})
planner = TensorAllocationPlanner(config={})

ugraph = parser.parse('tests/deep_mlp/simple_mnist.pb', output_nodes=['y_pred'])
planner.apply(ugraph)
plan = ugraph.attributes[planner.KWARGS_NAMESCOPE].plan
total_size = ugraph.attributes[planner.KWARGS_NAMESCOPE].total_size

for tensor_name, alloc in plan.items():
    print(f'{tensor_name!r}', f'memory offset: {alloc.offset_start}, size: {alloc.size}')
print('total memory usage:', total_size)