1 #!/usr/bin/python3
2"""Read intermediate tensors generated by DumpAllTensors activity
3
4Tools for reading/ parsing intermediate tensors.
5"""
6
7import argparse
8import numpy as np
9import os
10import pandas as pd
11import tensorflow as tf
12import matplotlib.pyplot as plt
13import json
14
15from matplotlib.pylab import *
16import matplotlib.animation as animation
17# Enable tensor.numpy()
18tf.compat.v1.enable_eager_execution()
19
20################################ ModelMetaDataManager ################################
21class ModelMetaDataManager(object):
22  """Maps model name in nnapi to its graph architecture with lazy initialization.
23
24  # Arguments
25    android_build_top: the root directory of android source tree
26    dump_dir: directory containing intermediate tensors pulled from device
27    tflite_model_json_path: directory containing intermediate json output of
28    model visualization tool (third_party/tensorflow/lite/tools:visualize)
29    The json output path from the tool is always /tmp.
30  """
31
32  class ModelMetaData(object):
33    """Store graph information of a model."""
34
35    def __init__(self, tflite_model_json_path='/tmp'):
36      with open(tflite_model_json_path, 'rb') as f:
37        model_json = json.load(f)
38      self.operators = model_json['subgraphs'][0]['operators']
39      self.operator_codes = [item['builtin_code']\
40                            for item in model_json['operator_codes']]
41      self.output_meta_data = []
42      self.load_output_meta_data()
43
44    def load_output_meta_data(self):
45      for operator in self.operators:
46        data = {}
47        # Each operator can only have one output
48        assert(len(operator['outputs']) == 1)
49        data['output_tensor_index'] = operator['outputs'][0]
50        data['fused_activation_function'] = operator\
51          .get('builtin_options', {})\
52          .get('fused_activation_function', '')
53        data['operator_code'] = self.operator_codes[operator['opcode_index']]
54        self.output_meta_data.append(data)
55
56  def __init__(self, android_build_top, dump_dir, tflite_model_json_dir='/tmp'):
57    # key: nnapi model name, value: ModelMetaData
58    self.models = dict()
59    self.ANDROID_BUILD_TOP = android_build_top
60    self.TFLITE_MODEL_JSON_DIR = tflite_model_json_dir
61    self.DUMP_DIR = dump_dir
62    self.nnapi_to_tflite_name = dict()
63    self.tflite_to_nnapi_name = dict()
64    self.__load_mobilenet_topk_aosp__()
65    self.model_names = sorted(os.listdir(dump_dir))
66
67  def __load_mobilenet_topk_aosp__(self):
68    """Load information about tflite and nnapi model names."""
69    json_path = '{}/{}'.format(
70        self.ANDROID_BUILD_TOP,
71        'test/mlts/models/assets/models_list/mobilenet_topk_aosp.json')
72    with open(json_path, 'rb') as f:
73      topk_aosp = json.load(f)
74    for model in topk_aosp['models']:
75      self.nnapi_to_tflite_name[model['name']] = model['modelFile']
76      self.tflite_to_nnapi_name[model['modelFile']] = model['name']
77
78  def __get_model_json_path__(self, tflite_model_name):
79    """Return tflite model jason path."""
80    json_path = '{}/{}.json'.format(self.TFLITE_MODEL_JSON_DIR, tflite_model_name)
81    return json_path
82
83  def __load_model__(self, tflite_model_name):
84    """Initialize a ModelMetaData for this model."""
85    model = self.ModelMetaData(self.__get_model_json_path__(tflite_model_name))
86    nnapi_model_name = self.model_name_tflite_to_nnapi(tflite_model_name)
87    self.models[nnapi_model_name] = model
88
89  def model_name_nnapi_to_tflite(self, nnapi_model_name):
90    return self.nnapi_to_tflite_name.get(nnapi_model_name, nnapi_model_name)
91
92  def model_name_tflite_to_nnapi(self, tflite_model_name):
93    return self.tflite_to_nnapi_name.get(tflite_model_name, tflite_model_name)
94
95  def get_model_meta_data(self, nnapi_model_name):
96    """Retrieve the ModelMetaData with lazy initialization."""
97    tflite_model_name = self.model_name_nnapi_to_tflite(nnapi_model_name)
98    if nnapi_model_name not in self.models:
99      self.__load_model__(tflite_model_name)
100    return self.models[nnapi_model_name]
101
102  def generate_animation_html(self, output_file_path, model_names=None):
103    model_names = self.model_names if model_names is None else model_names
104    html_data = ''
105    for model_name in model_names:
106      print('processing', model_name)
107      html_data += '<h3>{}</h3>'.format(model_name)
108      model_data = ModelData(nnapi_model_name=model_name, manager=self)
109      ani = model_data.gen_error_hist_animation()
110      html_data += ani.to_jshtml()
111    with open(output_file_path, 'w') as f:
112      f.write(html_data)
113
114
115################################ TensorDict ################################
116class TensorDict(dict):
117  """A class to store cpu and nnapi tensors.
118
119  # Arguments
120    model_dir: directory containing intermediate tensors pulled from device
121  """
122  def __init__(self, model_dir):
123    super().__init__()
124    for useNNAPIDir in ['cpu', 'nnapi']:
125      dir_path = model_dir + useNNAPIDir + "/"
126      self[useNNAPIDir] = self.read_tensors_from_dir(dir_path)
127    self.tensor_sanity_check()
128    self.max_absolute_diff, self.min_absolute_diff = 0.0, 0.0
129    self.max_relative_diff, self.min_relative_diff = 0.0, 0.0
130    self.layers = sorted(self['cpu'].keys())
131    self.calc_range()
132
133  def bytes_to_numpy_tensor(self, file_path):
134    tensor_type = tf.int8 if 'quant' in file_path else tf.float32
135    with open(file_path, mode='rb') as f:
136      tensor_bytes = f.read()
137      tensor = tf.decode_raw(input_bytes=tensor_bytes, out_type=tensor_type)
138    if np.isnan(np.sum(tensor)):
139      print('WARNING: tensor contains inf or nan')
140    return tensor.numpy()
141
142  def read_tensors_from_dir(self, dir_path):
143    tensor_dict = dict()
144    for tensor_file in os.listdir(dir_path):
145      tensor = self.bytes_to_numpy_tensor(dir_path + tensor_file)
146      tensor_dict[tensor_file] = tensor
147    return tensor_dict
148
149  def tensor_sanity_check(self):
150    # Make sure the cpu tensors and nnapi tensors have the same outputs
151    assert(len(self['cpu']) == len(self['nnapi']))
152    key_diff = set(self['cpu'].keys()) - set(self['nnapi'].keys())
153    assert(len(key_diff) == 0)
154    print('Tensor sanity check passed')
155
156  def calc_range(self):
157    for layer in self.layers:
158      diff = self.calc_diff(layer, relative_error=False)
159      # update absolute max, min
160      self.max_absolute_diff = max(self.max_absolute_diff, np.max(diff))
161      self.min_absolute_diff = min(self.min_absolute_diff, np.min(diff))
162      self.absolute_range = max(abs(self.min_absolute_diff),
163                                abs(self.max_absolute_diff))
164
165  def calc_diff(self, layer, relative_error=True):
166    cpu_tensor = self['cpu'][layer]
167    nnapi_tensor = self['nnapi'][layer]
168    assert(cpu_tensor.shape == nnapi_tensor.shape)
169    diff = cpu_tensor - nnapi_tensor
170    if not relative_error:
171      return diff
172    diff = diff.astype(float)
173    cpu_tensor = cpu_tensor.astype(float)
174    max_cpu_nnapi_tensor = np.maximum(np.abs(cpu_tensor), np.abs(nnapi_tensor))
175    relative_diff = np.divide(diff, max_cpu_nnapi_tensor, out=np.zeros_like(diff),\
176                              where=max_cpu_nnapi_tensor>0)
177    relative_diff[relative_diff>1] = 1.0
178    relative_diff[relative_diff<-1] = -1.0
179    return relative_diff
180
181  def gen_tensor_diff_stats(self, relative_error=True, return_df=True, plot_diff=False):
182    stats = []
183    for layer in self.layers:
184      diff = self.calc_diff(layer, relative_error)
185      if plot_diff:
186        self.plot_tensor_diff(diff)
187      if return_df:
188        stats.append({
189          'layer': layer,
190          'min': np.min(diff),
191          'max': np.max(diff),
192          'mean': np.mean(diff),
193          'median': np.median(diff)
194        })
195    if return_df:
196      return pd.DataFrame(stats)
197
198  def plot_tensor_diff(diff):
199    plt.figure()
200    plt.hist(diff, bins=50, log=True)
201    plt.plot()
202
203
204################################ Model Data ################################
205
206class ModelData(object):
207  """A class to store all relevant inormation of a model.
208
209  # Arguments
210    nnapi_model_name: the name of the model
211    manager: ModelMetaDataManager
212  """
213  def __init__(self, nnapi_model_name, manager):
214    self.nnapi_model_name = nnapi_model_name
215    self.manager = manager
216    self.model_dir = self.get_target_model_dir(manager.DUMP_DIR, nnapi_model_name)
217    self.tensor_dict = TensorDict(self.model_dir)
218    self.mmd = manager.get_model_meta_data(nnapi_model_name)
219    self.stats = self.tensor_dict.gen_tensor_diff_stats(relative_error=True,
220                                                        return_df=True)
221    self.layers = sorted(self.tensor_dict['cpu'].keys())
222
223  def get_target_model_dir(self, dump_dir, target_model_name):
224    target_model_dir = dump_dir + target_model_name + "/"
225    return target_model_dir
226
227  def updateData(self, i, fig, ax1, ax2, bins=50):
228    operation = self.mmd.output_meta_data[i % len(self.mmd.output_meta_data)]['operator_code']
229    layer = self.layers[i]
230    subtitle = fig.suptitle('{} | {}\n{}'
231                      .format(self.nnapi_model_name, layer, operation),
232                      fontsize='x-large')
233    for ax in (ax1, ax2):
234        ax.clear()
235    ax1.set_title('Relative Error')
236    ax2.set_title('Absolute Error')
237    ax1.hist(self.tensor_dict.calc_diff(layer, relative_error=True), bins=bins,
238             range=(-1, 1), log=True)
239    absolute_range = self.tensor_dict.absolute_range
240    ax2.hist(self.tensor_dict.calc_diff(layer, relative_error=False), bins=bins,
241             range=(-absolute_range, absolute_range), log=True)
242
243  def gen_error_hist_animation(self):
244    # For fast testing, add [:10] to the end of next line
245    layers = self.layers
246    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12,9))
247    ani = animation.FuncAnimation(fig, self.updateData, len(layers),
248                                  fargs=(fig, ax1, ax2),
249                                  interval=200, repeat=False)
250    # close before return to avoid dangling plot
251    plt.close()
252    return ani
253
254  def plot_error_heatmap(self, target_layer, length=1):
255    target_diff = self.tensor_dict['cpu'][target_layer] - \
256                  self.tensor_dict['nnapi'][target_layer]
257    width = int(len(target_diff)/ length)
258    reshaped_target_diff = target_diff[:length * width].reshape(length, width)
259    fig, ax = subplots(figsize=(18, 5))
260    plt.title('Heat Map of Error between CPU and NNAPI')
261    plt.imshow(reshaped_target_diff, cmap='hot', interpolation='nearest')
262    plt.colorbar()
263    plt.show()
264
265
266################################NumpyEncoder ################################
267
268class NumpyEncoder(json.JSONEncoder):
269  """Enable numpy array serilization in a dictionary.
270
271  # Usage:
272    a = np.array([[1, 2, 3], [4, 5, 6]])
273    json.dumps({'a': a, 'aa': [2, (2, 3, 4), a], 'bb': [2]}, cls=NumpyEncoder)
274  """
275  def default(self, obj):
276      if isinstance(obj, np.ndarray):
277          return obj.tolist()
278      return json.JSONEncoder.default(self, obj)
279
280def main(android_build_top, dump_dir, model_name):
281  manager = ModelMetaDataManager(
282    android_build_top,
283    dump_dir,
284    tflite_model_json_dir='/tmp')
285  model_data = ModelData(nnapi_model_name=model_name, manager=manager)
286  print(model_data.tensor_dict)
287
288if __name__ == '__main__':
289  # Example usage
290  # python tensor_utils.py ~/android/master/ ~/android/master/intermediate/ tts_float
291  parser = argparse.ArgumentParser(description='Utilities for parsing intermediate tensors.')
292  parser.add_argument('android_build_top', help='Your Android build top path')
293  parser.add_argument('dump_dir', help='The dump dir pulled from the device')
294  parser.add_argument('model_name', help='NNAPI model name')
295  args = parser.parse_args()
296  main(args.android_build_top, args.dump_dir, args.model_name)