【ATU Book-i.MX8系列 - TFLite 進階】 模組轉換(三)

$ git clone https://github.com/google/flatbuffers.git 
$ cd /root/flatbuffers
$ cmake -G "Unix Makefiles"
$ make

$ git clone https://github.com/google/mediapipe
$ cd /root/mediapipe
$ cp mediapipe/modules/palm_detection/palm_detection.tflite /root/flatbuffers/palm_detection.tflite  # 移動至flatbuffers目錄下​

$ git clone https://github.com/tensorflow/tensorflow
$ cd /root/tensorflow
$ cp tensorflow/lite/schema/schema.fbs /root/flatbuffers/schema.fbs  # 移動至flatbuffers目錄下

# (1) 載入必要函式庫
import os
import numpy as np
import json
import tensorflow as tf
import shutil
from pathlib import Path

# (2) API config
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
schema = "schema.fbs"
binary = "./flatc"
model_path = "palm_detection.tflite"
output_pb_path = "palm_detection.pb"
output_savedmodel_path = "saved_model"
model_json_path = "palm_detection.json"
output_node_names = [ 'classificators' , 'regressors' ]

#( 3) API : 共三個
# 產生架構檔(JSON)
def gen_model_json():
  cmd = (binary + " -t --strict-json --defaults-json -o . {schema} -- {input}".format(input=model_path, schema=schema))
  print("output json command =", cmd)
  os.system(cmd)

# 解析架構
def parse_json():
    j = json.load(open(model_json_path))
    op_types = [ v['builtin_code’] for v in j[ 'operator_codes'] ]
    ops = j['subgraphs'][0]['operators'] # subgraphs => 分為 tensors / inputs / outputs / operators
    return ops, op_types

# 描述架構圖
# make_graph 函式內需要撰寫 tflite 內所有用到 架構層 layer 
# 因此非常龐大 這裡僅介紹 conv2 的描述，完整代碼請至 “該連結查看”
def make_graph(ops, op_types, interpreter):
    tensors = {}
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    for input_detail in input_details:
        tensors[input_detail['index']] = tf.compat.v1.placeholder(
            dtype=input_detail['dtype'], shape=input_detail['shape'], name=input_detail['name'])

    for index, op in enumerate(ops):
        print('op: ', op)
        op_type = op_types[op['opcode_index']]
        if op_type == 'CONV_2D':
            input_tensor = tensors[op['inputs'][0]]
            weights_detail = interpreter._get_tensor_details(op['inputs'][1])
            bias_detail = interpreter._get_tensor_details(op['inputs'][2])
            output_detail = interpreter._get_tensor_details(op['outputs'][0])
            weights_array = interpreter.get_tensor(weights_detail['index'])
            weights_array = np.transpose(weights_array, (1, 2, 3, 0))
            bias_array = interpreter.get_tensor(bias_detail['index'])
            weights = tf.compat.v1.Variable(weights_array, name=weights_detail['name'])
            bias = tf.compat.v1.Variable(bias_array, name=bias_detail['name'])
            options = op['builtin_options']
            output_tensor = tf.compat.v1.nn.conv2d(
                input_tensor,
                weights,
                strides=[1, options['stride_h'], options['stride_w'], 1],
                padding=options['padding'],
                dilations=[1, options['dilation_h_factor'],options['dilation_w_factor'], 1],
                name=output_detail['name'] + '/conv2d')
            output_tensor = tf.compat.v1.add(output_tensor, bias, name=output_detail['name'])

# 將 tf lite 架構轉成 json 格式
tf.compat.v1.disable_eager_execution()
gen_model_json()
ops, op_types = parse_json() 

# 調用 tflite 取得資訊
interpreter = Interpreter(model_path)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(input_details)
print(output_details)

# 將 tf lite 重新拆解回 savedmodel
# 若欲將 tflite 轉成全整數的格式，請查閱並結合上一小節方式。
make_graph(ops, op_types, interpreter) 
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
graph = tf.compat.v1.get_default_graph()
with tf.compat.v1.Session(config=config, graph=graph) as sess:
  sess.run(tf.compat.v1.global_variables_initializer())
  graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(sess=sess,\
                                                  input_graph_def=graph.as_graph_def(),\
                                                  output_node_names=output_node_names)
  with tf.io.gfile.GFile(output_pb_path, 'wb') as f:
      f.write(graph_def.SerializeToString())
      shutil.rmtree('saved_model', ignore_errors=True)
      tf.compat.v1.saved_model.simple_save(sess,
            output_savedmodel_path,
            inputs={'input': graph.get_tensor_by_name('input:0')},
            outputs={
                'classificators': graph.get_tensor_by_name('classificators:0'),
                'regressors': graph.get_tensor_by_name('regressors:0')
            })

# 因擔心代碼過於冗長，故此代碼轉成 tflite 格式為主。 
import tensorflow as tf
converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model( 'saved_model' )
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.compat.v1.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type  = tf.compat.v1.uint8
converter.inference_output_type = tf.compat.v1.uint8
converter.representative_dataset = representative_dataset_gen # representative_dataset 請參閱上述的量化方式
tflite_model = converter.convert()
with tf.io.gfile.GFile( "/root/palm_detection_uint8.tflite" , 'wb') as f:
   f.write(tflite_model)
print("Integer Quantization complete! ")