diff options
Diffstat (limited to 'ethosu/vela/tensor_allocation.py')
| -rw-r--r-- | ethosu/vela/tensor_allocation.py | 14 |
1 files changed, 9 insertions, 5 deletions
diff --git a/ethosu/vela/tensor_allocation.py b/ethosu/vela/tensor_allocation.py index 7f66579e..1e5eb852 100644 --- a/ethosu/vela/tensor_allocation.py +++ b/ethosu/vela/tensor_allocation.py @@ -202,8 +202,11 @@ def allocate_tensors( else: sg.memory_used_per_type[mem_type] += total_sz - nng.total_size[mem_area] = nng.total_size.get(mem_area, 0) + sum(tens.storage_size() for tens in lrs.ranges) - nng.total_elements[mem_area] = nng.total_elements.get(mem_area, 0) + sum(tens.elements() for tens in lrs.ranges) + if mem_area == arch.fast_storage_mem_area: + for tens in lrs.ranges: + if tens.purpose == TensorPurpose.Weights: + nng.total_compressed_weights += tens.storage_size() + nng.total_original_weights += tens.elements() * tens.element_size() print_allocation(lrs, mem_area, mem_type_set, sg, verbose_allocation) @@ -214,9 +217,10 @@ def allocate_tensors( if sg == nng.get_root_subgraph(): nng.memory_used = sg.memory_used - for mem_area in nng.total_elements.keys(): + if mem_area == arch.fast_storage_mem_area: try: - nng.bits_per_element[mem_area] = nng.total_size[mem_area] * 8 / nng.total_elements[mem_area] + nng.weights_compression_ratio = nng.total_compressed_weights / nng.total_original_weights except ZeroDivisionError: - nng.bits_per_element[mem_area] = 0.0 + nng.weights_compression_ratio = 0.0 + return True |