From 535b36be0c2630d10affbd0dc560e463069736bb Mon Sep 17 00:00:00 2001
From: Yaman Umuroglu <maltanar@gmail.com>
Date: Tue, 11 Feb 2020 18:50:08 +0100
Subject: [PATCH] [ModelWrapper] add helper method to get ValueInfo
---
src/finn/core/modelwrapper.py | 13 +++++++++++++
1 file changed, 13 insertions(+)
diff --git a/src/finn/core/modelwrapper.py b/src/finn/core/modelwrapper.py
index 1d5d3a978..473a33c47 100644
--- a/src/finn/core/modelwrapper.py
+++ b/src/finn/core/modelwrapper.py
@@ -115,6 +115,19 @@ class ModelWrapper:
qa.quant_parameter_tensor_names.append(dt)
qnt_annotations.append(qa)
+ def get_tensor_valueinfo(self, tensor_name):
+ """Returns ValueInfoProto of tensor with given name, if it has one."""
+ graph = self._model_proto.graph
+ vi_names = [(x.name, x) for x in graph.input]
+ vi_names += [(x.name, x) for x in graph.output]
+ vi_names += [(x.name, x) for x in graph.value_info]
+ try:
+ vi_ind = [x[0] for x in vi_names].index(tensor_name)
+ vi = vi_names[vi_ind][1]
+ return vi
+ except ValueError:
+ return None
+
def get_tensor_shape(self, tensor_name):
"""Returns the shape of tensor with given name, if it has ValueInfoProto."""
graph = self._model_proto.graph
--
GitLab