diff options
| -rw-r--r-- | tensorflow/contrib/lite/python/convert.py | 4 | ||||
| -rw-r--r-- | tensorflow/contrib/lite/toco/tflite/export.cc | 14 |
2 files changed, 11 insertions, 7 deletions
diff --git a/tensorflow/contrib/lite/python/convert.py b/tensorflow/contrib/lite/python/convert.py index c4200c879b..c0926d2f33 100644 --- a/tensorflow/contrib/lite/python/convert.py +++ b/tensorflow/contrib/lite/python/convert.py @@ -118,7 +118,8 @@ def toco_convert(input_data, input_format=lite_constants.TENSORFLOW_GRAPHDEF, output_format=lite_constants.TFLITE, quantized_input_stats=None, - drop_control_dependency=True): + drop_control_dependency=True, + allow_custom_ops=False): """Convert a model using TOCO from `input_format` to `output_format`. Typically this is to convert from TensorFlow GraphDef to TFLite, in which @@ -154,6 +155,7 @@ def toco_convert(input_data, toco.drop_control_dependency = drop_control_dependency model = _model_flags_pb2.ModelFlags() toco.inference_type = inference_type + toco.allow_custom_ops = allow_custom_ops for idx, input_tensor in enumerate(input_tensors): if input_tensor.dtype == _dtypes.float32: tflite_input_type = lite_constants.FLOAT diff --git a/tensorflow/contrib/lite/toco/tflite/export.cc b/tensorflow/contrib/lite/toco/tflite/export.cc index 335b496dcc..a4c0b2d16e 100644 --- a/tensorflow/contrib/lite/toco/tflite/export.cc +++ b/tensorflow/contrib/lite/toco/tflite/export.cc @@ -312,12 +312,14 @@ void Export(const Model& model, bool allow_custom_ops, error_summary.erase(fake_quant_operation_name); } if (!allow_custom_ops && !error_summary.empty()) { - LOG(QFATAL) << "Some of the operators in the model are not supported by " - "the standard TensorFlow Lite runtime. If you have a custom " - "implementation for them you can disable this error with " - "--allow_custom_ops. Here is a list of operators for which " - "you will need custom implementations: " - << absl::StrJoin(error_summary, ", ") << "."; + LOG(QFATAL) + << "Some of the operators in the model are not supported by " + "the standard TensorFlow Lite runtime. If you have a custom " + "implementation for them you can disable this error with " + "--allow_custom_ops, or by setting allow_custom_ops=True " + "when calling tf.contrib.lite.toco_convert(). Here is a list " + "of operators for which you will need custom implementations: " + << absl::StrJoin(error_summary, ", ") << "."; } auto ops = |