def collect(self):
"""Callback function for collecting layer output NDArrays."""
# name = py_str(name)
# if name not in self.include_layer:
# return
# handle = ctypes.cast(arr, NDArrayHandle)
# arr = NDArray(handle, writable=False).copyto(cpu()).asnumpy()
for name in self.layer_tensor:
if name not in self.include_layer:
continue
for arr in self.layer_tensor[name]:
if self.logger:
self.logger.debug(
"Collecting layer %s histogram of shape %s" %
(name, arr.shape))
min_range = np.min(arr)
max_range = np.max(arr)
th = max(abs(min_range), abs(max_range))
if name in self.hist_dict:
self.hist_dict[name] = combine_histogram(
self.hist_dict[name], arr)
else:
hist, hist_edges = np.histogram(arr,
bins=self.num_bins,
range=(-th, th))
self.hist_dict[name] = (hist, hist_edges, min_range,
max_range, th)
def _generate_calibration_data(self,
tmp_path,
output_data,
enable_kl_algo=False):
tmp_dump_file = os.path.join(os.path.dirname(self.output_graph),
'requant_min_max.log')
self.logger.debug(
"Generating calibration data and saving to {}".format(
tmp_dump_file))
model = TensorflowModel(tmp_path,
self._tmp_model.framework_specific_info,
**self._tmp_model.kwargs)
with CaptureOutputToFile(tmp_dump_file):
self._inference(model)
with open(tmp_dump_file, errors='ignore') as f:
output_data.extend(f.readlines())
for line in output_data:
if enable_kl_algo and line.rsplit(':')[0] in self._kl_keys:
fp32_data = get_all_fp32_data(line.rsplit(':')[-1])
key = self._print_node_mapping[line[1:].split('__print')
[0]] + '_eightbit_requant_range'
if key not in self._kl_op_dict:
self._kl_op_dict[key] = get_tensor_histogram(fp32_data)
else:
self._kl_op_dict[key] = combine_histogram(
self._kl_op_dict[key], fp32_data)