Create TOCO options to process a model. Args: data_types: input and inference types used by TOCO. input_arrays: names of the input tensors output_arrays: name of the output tensors shapes: shapes of the input tensors extra_toco_options: additional toco options Returns:
(data_types,
input_arrays,
output_arrays,
shapes,
extra_toco_options=None)
| 28 | |
| 29 | |
| 30 | def toco_options(data_types, |
| 31 | input_arrays, |
| 32 | output_arrays, |
| 33 | shapes, |
| 34 | extra_toco_options=None): |
| 35 | """Create TOCO options to process a model. |
| 36 | |
| 37 | Args: |
| 38 | data_types: input and inference types used by TOCO. |
| 39 | input_arrays: names of the input tensors |
| 40 | output_arrays: name of the output tensors |
| 41 | shapes: shapes of the input tensors |
| 42 | extra_toco_options: additional toco options |
| 43 | |
| 44 | Returns: |
| 45 | the options in a string. |
| 46 | """ |
| 47 | if extra_toco_options is None: |
| 48 | extra_toco_options = generate_examples_lib.ExtraTocoOptions() |
| 49 | |
| 50 | shape_str = ":".join([",".join(str(y) for y in x) for x in shapes if x]) |
| 51 | inference_type = "FLOAT" |
| 52 | # TODO(ahentz): if we get multi-input quantization to work we need this |
| 53 | # to change |
| 54 | if data_types[0] == "QUANTIZED_UINT8": |
| 55 | inference_type = "QUANTIZED_UINT8" |
| 56 | s = (" --input_data_types=%s" % ",".join(data_types) + |
| 57 | " --inference_type=%s" % inference_type + |
| 58 | " --input_format=TENSORFLOW_GRAPHDEF" + " --output_format=TFLITE" + |
| 59 | " --input_arrays=%s" % ",".join(input_arrays) + |
| 60 | " --output_arrays=%s" % ",".join(output_arrays)) |
| 61 | if shape_str: |
| 62 | s += (" --input_shapes=%s" % shape_str) |
| 63 | if extra_toco_options.drop_control_dependency: |
| 64 | s += " --drop_control_dependency" |
| 65 | if extra_toco_options.allow_custom_ops: |
| 66 | s += " --allow_custom_ops" |
| 67 | if extra_toco_options.rnn_states: |
| 68 | s += (" --rnn_states='" + extra_toco_options.rnn_states + "'") |
| 69 | if extra_toco_options.split_tflite_lstm_inputs is not None: |
| 70 | if extra_toco_options.split_tflite_lstm_inputs: |
| 71 | s += " --split_tflite_lstm_inputs=true" |
| 72 | else: |
| 73 | s += " --split_tflite_lstm_inputs=false" |
| 74 | return s |
| 75 | |
| 76 | |
| 77 | def toco_convert(options, graph_def, input_tensors, output_tensors, **kwargs): |