def _move_quant_attributes_into_annotations(model):
"""Move quantization info in attributes into quantization_annotation"""
if onnx is None:
raise ModuleNotFoundError("Installation of ONNX is required.")
model = copy.deepcopy(model)
qaname = "finn_datatype"
for n in model.graph.node:
for a in n.attribute:
mark_for_removal = False
if a.name == "weight_qnt":
# assume second input is weight, make sure it has an initializer
w_tensor_name = n.input[1]
assert w_tensor_name in [
x.name for x in model.graph.initializer
]
tq = onnx.StringStringEntryProto(key=qaname, value=a.s)
ta = onnx.TensorAnnotation(tensor_name=w_tensor_name,
quant_parameter_tensor_names=[tq])
model.graph.quantization_annotation.append(ta)
mark_for_removal = True
elif a.name == "activation_qnt":
a_tensor_name = n.output[0]
tq = onnx.StringStringEntryProto(key=qaname, value=a.s)
ta = onnx.TensorAnnotation(tensor_name=a_tensor_name,
quant_parameter_tensor_names=[tq])
model.graph.quantization_annotation.append(ta)
mark_for_removal = True
if mark_for_removal:
n.attribute.remove(a)
return model
def set_tensor_layout(self, tensor_name, data_layout):
"""Sets the data layout annotation of tensor with given name. See
get_tensor_layout for examples."""
tensor_shape = self.get_tensor_shape(tensor_name)
assert type(data_layout) == list, "data_layout must be a list"
if tensor_shape is not None:
assert len(tensor_shape) == len(
data_layout
), """Mismatch between number
of dimensions of tensor shape and data layout annotation."""
graph = self._model_proto.graph
qnt_annotations = graph.quantization_annotation
ret = util.get_by_name(qnt_annotations, tensor_name, "tensor_name")
if ret is not None:
ret_tl = util.get_by_name(
ret.quant_parameter_tensor_names, "tensor_layout", "key"
)
if ret_tl is not None:
ret_tl.value = str(data_layout)
else:
tl = onnx.StringStringEntryProto()
tl.key = "tensor_layout"
tl.value = str(data_layout)
ret.quant_parameter_tensor_names.append(tl)
else:
qa = onnx.TensorAnnotation()
dt = onnx.StringStringEntryProto()
dt.key = "tensor_layout"
dt.value = str(data_layout)
qa.tensor_name = tensor_name
qa.quant_parameter_tensor_names.append(dt)
qnt_annotations.append(qa)
def set_tensor_datatype(self, tensor_name, datatype):
"""Sets the FINN DataType of tensor with given name."""
graph = self._model_proto.graph
qnt_annotations = graph.quantization_annotation
ret = util.get_by_name(qnt_annotations, tensor_name, "tensor_name")
if ret is not None:
ret_dt = util.get_by_name(ret.quant_parameter_tensor_names,
"finn_datatype", "key")
if ret_dt is not None:
if datatype is None:
ret_dt.Clear()
else:
ret_dt.value = datatype.name
elif datatype is not None:
dt = onnx.StringStringEntryProto()
dt.key = "finn_datatype"
dt.value = datatype.name
ret.quant_parameter_tensor_names.append(dt)
elif datatype is not None:
qa = onnx.TensorAnnotation()
dt = onnx.StringStringEntryProto()
dt.key = "finn_datatype"
dt.value = datatype.name
qa.tensor_name = tensor_name
qa.quant_parameter_tensor_names.append(dt)
qnt_annotations.append(qa)
def set_tensor_sparsity(self, tensor_name, sparsity_dict):
"""Sets the sparsity annotation of a tensor with given name."""
graph = self._model_proto.graph
qnt_annotations = graph.quantization_annotation
ret = util.get_by_name(qnt_annotations, tensor_name, "tensor_name")
if ret is not None:
ret_ts = util.get_by_name(ret.quant_parameter_tensor_names,
"tensor_sparsity", "key")
if ret_ts is not None:
ret_ts.value = str(sparsity_dict)
else:
ts = onnx.StringStringEntryProto()
ts.key = "tensor_sparsity"
ts.value = str(sparsity_dict)
ret.quant_parameter_tensor_names.append(ts)
else:
qa = onnx.TensorAnnotation()
dt = onnx.StringStringEntryProto()
dt.key = "tensor_sparsity"
dt.value = str(sparsity_dict)
qa.tensor_name = tensor_name
qa.quant_parameter_tensor_names.append(dt)
qnt_annotations.append(qa)