tflite の Sparseテンソル Densify を逆量子化して Float32 のクリーンなモデルを錬成する

tensorflow-onnx が Densify という Sparse モデルの変換に対応していなかったので自作の逆量子化ツールへ実装した。ちなみに、Densify は interpreter.get_tensor を実行したときに Segmentation Fault で TensorFlow Lite のランタイムがAbortするため、普通には情報が読み取れない状況だった。よって、.tflite のバイナリを一時的にJSONへパースし、重みを含めて全ての情報を可読な状態へ加工してから Sparseテンソル を Denseテンソル へ逆変換した。

https://github.com/PINTO0309/tflite2tensorflow

Densify はモデルのサイズを圧縮するため、ゼロに近似する値を全て FlatBuffer から削除して記録しているため、ランタイムでそのまま読み込もうとすると構造上定義されているテンソルのサイズと .tflite に記録されているバイトサイズが一致せず Abort する模様。

JSONへパースした重みを Float16 のバイト配列へ変換し、Sparse化されて削除されたゼロを補間していく。

elif op_type == 'DENSIFY':
    json_tensor_info = searh_json_tensor_detail(
	interpreter._get_tensor_details(
	    op['outputs'][0]
	)['name'][:-1]
    )
    output_detail = interpreter._get_tensor_details(
	op['outputs'][0]
    )

    shape = json_tensor_info['shape']
    dtype = cast_type_np[json_tensor_info['type']]
    dim_metadata = json_tensor_info['sparsity']['dim_metadata']
    traversal_order = json_tensor_info['sparsity']['traversal_order']
    array_segments = None
    array_indices = None
    for dict_data in dim_metadata:
	if 'format' in dict_data:
	    if dict_data['format'] == 'SPARSE_CSR':
		array_segments = dict_data['array_segments']['values']
		array_indices = dict_data['array_indices']['values']

    denj = full_json['buffers'][output_detail['index']]['data']
    b = np.array(denj).astype(np.uint8).tobytes()

    dense_list = []
    for i in range(len(b))[::2]:
	dense_list.append(struct.unpack_from('<e', bytes(b[i:i+2]))[0])

    starting_point = 0
    dense_shape = np.asarray(shape)
    groups = dense_shape[-1]
    total_number_of_elements = dense_shape[0]
    for i in dense_shape[1:]:
	total_number_of_elements *= i

    densify_values = np.zeros((total_number_of_elements))
    sidx = 0
    aidx = 0
    didx = 0
    addition = 0

    for idx in range(total_number_of_elements):
	if array_segments[sidx] == aidx:
	    addition = sidx * groups
	    sidx += 1
	if array_indices[aidx] + addition == idx:
	    densify_values[idx] = dense_list[didx]
	    didx += 1
	    if aidx < len(array_indices) - 1:
		aidx += 1

    densify_values = densify_values.reshape(dense_shape)
    output_tensor = densify_values.astype(dtype)

    tensors[output_detail['index']] = output_tensor

tflite2tensorflow を使えば上記の実装を意識しなくてもモデルを逆変換して saved_model や Float32 の tflite へ加工することができるので特に難しくはない。とりあえず実装上の覚書。