# /path/to/tensorflow: tensorflow root path when use tflite library
set(USE_TENSORFLOW_PATH none)
+# Possible values:
+# - OFF: disable tflite support for edgetpu
+# - /path/to/edgetpu: use specific path to edgetpu library
+set(USE_EDGETPU OFF)
+
# Whether use CuDNN
set(USE_CUDNN OFF)
list(APPEND RUNTIME_SRCS ${TFLITE_CONTRIB_SRC})
include_directories(${USE_TENSORFLOW_PATH})
+ # Additional EdgeTPU libs
+ if (NOT USE_EDGETPU STREQUAL "OFF")
+ message(STATUS "Build with contrib.edgetpu")
+ file(GLOB EDGETPU_CONTRIB_SRC src/runtime/contrib/edgetpu/*.cc)
+ list(APPEND RUNTIME_SRCS ${EDGETPU_CONTRIB_SRC})
+ include_directories(${USE_EDGETPU}/libedgetpu)
+ list(APPEND TVM_RUNTIME_LINKER_LIBS ${USE_EDGETPU}/libedgetpu/direct/aarch64/libedgetpu.so.1)
+ endif()
+
if (USE_TFLITE STREQUAL "ON")
set(USE_TFLITE ${USE_TENSORFLOW_PATH}/tensorflow/lite/tools/make/gen/*/lib)
endif()
from .._ffi.function import get_global_func
from ..rpc import base as rpc_base
-def create(tflite_model_bytes, ctx):
+def create(tflite_model_bytes, ctx, runtime_target='cpu'):
"""Create a runtime executor module given a tflite model and context.
Parameters
----------
ctx : TVMContext
The context to deploy the module. It can be local or remote when there
is only one TVMContext.
+ runtime_target: str
+ Execution target of TFLite runtime: either `cpu` or `edge_tpu`.
Returns
-------
tflite_runtime : TFLiteModule
Runtime tflite module that can be used to execute the tflite model.
"""
device_type = ctx.device_type
+
+ if runtime_target == 'edge_tpu':
+ runtime_func = "tvm.edgetpu_runtime.create"
+ else:
+ runtime_func = "tvm.tflite_runtime.create"
+
if device_type >= rpc_base.RPC_SESS_MASK:
- fcreate = ctx._rpc_sess.get_function("tvm.tflite_runtime.create")
- return TFLiteModule(fcreate(bytearray(tflite_model_bytes), ctx))
- fcreate = get_global_func("tvm.tflite_runtime.create")
+ fcreate = ctx._rpc_sess.get_function(runtime_func)
+ else:
+ fcreate = get_global_func(runtime_func)
+
return TFLiteModule(fcreate(bytearray(tflite_model_bytes), ctx))
Parameters
----------
module : Module
- The interal tvm module that holds the actual tflite functions.
+ The internal tvm module that holds the actual tflite functions.
Attributes
----------
module : Module
- The interal tvm module that holds the actual tflite functions.
+ The internal tvm module that holds the actual tflite functions.
"""
def __init__(self, module):
self._set_input = module["set_input"]
self._invoke = module["invoke"]
self._get_output = module["get_output"]
- self._allocate_tensors = module["allocate_tensors"]
def set_input(self, index, value):
"""Set inputs to the module via kwargs
"""
self._invoke()
- def allocate_tensors(self):
- """Allocate space for all tensors.
- """
- self._allocate_tensors()
-
-
def get_output(self, index):
"""Get index-th output to out
--- /dev/null
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file edgetpu_runtime.cc
+ */
+#include <tvm/runtime/registry.h>
+#include <tensorflow/lite/interpreter.h>
+#include <tensorflow/lite/kernels/register.h>
+#include <tensorflow/lite/model.h>
+#include <edgetpu.h>
+
+
+#include "edgetpu_runtime.h"
+
+namespace tvm {
+namespace runtime {
+
+void EdgeTPURuntime::Init(const std::string& tflite_model_bytes,
+ TVMContext ctx) {
+ const char* buffer = tflite_model_bytes.c_str();
+ size_t buffer_size = tflite_model_bytes.size();
+ // Load compiled model as a FlatBufferModel
+ std::unique_ptr<tflite::FlatBufferModel> model =
+ tflite::FlatBufferModel::BuildFromBuffer(buffer, buffer_size);
+ // Build resolver
+ tflite::ops::builtin::BuiltinOpResolver resolver;
+ // Init EdgeTPUContext object
+ edgetpu_context_ = edgetpu::EdgeTpuManager::GetSingleton()->OpenDevice();
+ // Add custom edgetpu ops to resolver
+ resolver.AddCustom(edgetpu::kCustomOp, edgetpu::RegisterCustomOp());
+ // Build interpreter
+ TfLiteStatus status = tflite::InterpreterBuilder(*model, resolver)(&interpreter_);
+ CHECK_TFLITE_STATUS(status) << "Failed to build interpreter.";
+ // Bind EdgeTPU context with interpreter.
+ interpreter_->SetExternalContext(kTfLiteEdgeTpuContext, edgetpu_context_.get());
+ interpreter_->SetNumThreads(1);
+ // Allocate tensors
+ status = interpreter_->AllocateTensors();
+ CHECK_TFLITE_STATUS(status) << "Failed to allocate tensors.";
+
+ ctx_ = ctx;
+}
+
+Module EdgeTPURuntimeCreate(const std::string& tflite_model_bytes,
+ TVMContext ctx) {
+ auto exec = make_object<EdgeTPURuntime>();
+ exec->Init(tflite_model_bytes, ctx);
+ return Module(exec);
+}
+
+TVM_REGISTER_GLOBAL("tvm.edgetpu_runtime.create")
+ .set_body([](TVMArgs args, TVMRetValue* rv) {
+ *rv = EdgeTPURuntimeCreate(args[0], args[1]);
+ });
+} // namespace runtime
+} // namespace tvm
--- /dev/null
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \brief EdgeTPU runtime that can run tflite model compiled
+ * for EdgeTPU containing only tvm PackedFunc.
+ * \file edgetpu_runtime.h
+ */
+#ifndef TVM_RUNTIME_CONTRIB_EDGETPU_EDGETPU_RUNTIME_H_
+#define TVM_RUNTIME_CONTRIB_EDGETPU_EDGETPU_RUNTIME_H_
+
+#include <string>
+#include <memory>
+
+#include "../tflite/tflite_runtime.h"
+
+namespace tvm {
+namespace runtime {
+
+/*!
+ * \brief EdgeTPU runtime.
+ *
+ * This runtime can be accessed in various languages via
+ * the TVM runtime PackedFunc API.
+ */
+class EdgeTPURuntime : public TFLiteRuntime {
+ public:
+ /*!
+ * \return The type key of the executor.
+ */
+ const char* type_key() const final {
+ return "EdgeTPURuntime";
+ }
+
+ /*!
+ * \brief Initialize the edge TPU tflite runtime with tflite model and context.
+ * \param tflite_model_bytes The tflite model.
+ * \param ctx The context where the tflite model will be executed on.
+ */
+ void Init(const std::string& tflite_model_bytes,
+ TVMContext ctx);
+
+ private:
+ std::shared_ptr<edgetpu::EdgeTpuContext> edgetpu_context_;
+};
+
+} // namespace runtime
+} // namespace tvm
+
+#endif // TVM_RUNTIME_CONTRIB_EDGETPU_EDGETPU_RUNTIME_H_
* \file tflite_runtime.cc
*/
#include <tvm/runtime/registry.h>
-#include <tvm/dtype.h>
#include <tensorflow/lite/interpreter.h>
#include <tensorflow/lite/kernels/register.h>
#include <tensorflow/lite/model.h>
namespace runtime {
#define TVM_DTYPE_DISPATCH(type, DType, ...) \
- if (type == DataType::Float(64)) { \
+ if (type == DataType::Float(64)) { \
typedef double DType; \
{__VA_ARGS__} \
- } else if (type == DataType::Float(32)) { \
+ } else if (type == DataType::Float(32)) { \
typedef float DType; \
{__VA_ARGS__} \
- } else if (type == DataType::Float(16)) { \
+ } else if (type == DataType::Float(16)) { \
typedef uint16_t DType; \
{__VA_ARGS__} \
- } else if (type == DataType::Int(64)) { \
+ } else if (type == DataType::Int(64)) { \
typedef int64_t DType; \
{__VA_ARGS__} \
- } else if (type == DataType::Int(32)) { \
+ } else if (type == DataType::Int(32)) { \
typedef int32_t DType; \
{__VA_ARGS__} \
- } else if (type == DataType::Int(16)) { \
+ } else if (type == DataType::Int(16)) { \
typedef int16_t DType; \
{__VA_ARGS__} \
- } else if (type == DataType::Int(8)) { \
+ } else if (type == DataType::Int(8)) { \
typedef int8_t DType; \
{__VA_ARGS__} \
- } else if (type == DataType::UInt(64)) { \
+ } else if (type == DataType::UInt(64)) { \
typedef uint64_t DType; \
{__VA_ARGS__} \
- } else if (type == DataType::UInt(32)) { \
+ } else if (type == DataType::UInt(32)) { \
typedef uint32_t DType; \
{__VA_ARGS__} \
- } else if (type == DataType::UInt(16)) { \
+ } else if (type == DataType::UInt(16)) { \
typedef uint16_t DType; \
{__VA_ARGS__} \
- } else if (type == DataType::UInt(8)) { \
+ } else if (type == DataType::UInt(8)) { \
typedef uint8_t DType; \
{__VA_ARGS__} \
} else { \
case kTfLiteInt64:
return DataType::Int(64);
case kTfLiteInt16:
- returnDataType::Int(16);
+ return DataType::Int(16);
case kTfLiteInt8:
- returnDataType::Int(8);
+ return DataType::Int(8);
case kTfLiteUInt8:
return DataType::UInt(8);
case kTfLiteFloat16:
}
}
-
void TFLiteRuntime::Init(const std::string& tflite_model_bytes,
TVMContext ctx) {
const char* buffer = tflite_model_bytes.c_str();
std::unique_ptr<tflite::FlatBufferModel> model =
tflite::FlatBufferModel::BuildFromBuffer(buffer, buffer_size);
tflite::ops::builtin::BuiltinOpResolver resolver;
- tflite::InterpreterBuilder(*model, resolver)(&interpreter_);
- ctx_ = ctx;
-}
+ // Build interpreter
+ TfLiteStatus status = tflite::InterpreterBuilder(*model, resolver)(&interpreter_);
+ CHECK_TFLITE_STATUS(status) << "Failed to build interpreter.";
+ // Allocate tensors
+ status = interpreter_->AllocateTensors();
+ CHECK_TFLITE_STATUS(status) << "Failed to allocate tensors.";
-void TFLiteRuntime::AllocateTensors() {
- interpreter_->AllocateTensors();
+ ctx_ = ctx;
}
void TFLiteRuntime::Invoke() {
}
NDArray TFLiteRuntime::GetOutput(int index) const {
- TfLiteTensor* output = interpreter_->output_tensor(index);
+ TfLiteTensor* output = interpreter_->tensor(interpreter_->outputs()[index]);
DataType dtype = TfLiteDType2TVMDType(output->type);
TfLiteIntArray* dims = output->dims;
int64_t size = 1;
return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
this->Invoke();
});
- } else if (name == "allocate_tensors") {
- return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
- this->AllocateTensors();
- });
} else {
return PackedFunc();
}
namespace tvm {
namespace runtime {
+#define CHECK_TFLITE_STATUS(ret) CHECK_EQ(ret, kTfLiteOk)
/*!
* \brief Tflite runtime.
*
- * This runtime can be acccesibly in various language via
+ * This runtime can be accessed in various language via
* TVM runtime PackedFunc API.
*/
class TFLiteRuntime : public ModuleNode {
public:
/*!
- * \brief Get member function to front-end
+ * \brief Get member function to front-end.
* \param name The name of the function.
* \param sptr_to_self The pointer to the module node.
* \return The corresponding member function.
/*!
* \return The type key of the executor.
*/
- const char* type_key() const final {
+ const char* type_key() const {
return "TFLiteRuntime";
}
/*!
- * \brief Update allocations for all tenssors. This is relatively expensive.
- */
- void AllocateTensors();
- /*!
* \brief Invoke the internal tflite interpreter and run the whole model in
* dependency order.
*/
*/
NDArray GetOutput(int index) const;
- private:
+ // TFLite interpreter
std::unique_ptr<tflite::Interpreter> interpreter_;
+ // TVM context
TVMContext ctx_;
};
--- /dev/null
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import os
+import tvm
+import numpy as np
+from tvm import rpc
+from tvm.contrib import util, tflite_runtime
+# import tflite_runtime.interpreter as tflite
+
+
+def skipped_test_tflite_runtime():
+
+ def get_tflite_model_path(target_edgetpu):
+ # Return a path to the model
+ edgetpu_path = os.getenv('EDGETPU_PATH', "/home/mendel/edgetpu")
+ # Obtain mobilenet model from the edgetpu repo path
+ if target_edgetpu:
+ model_path = os.path.join(edgetpu_path, "test_data/mobilenet_v1_1.0_224_quant_edgetpu.tflite")
+ else:
+ model_path = os.path.join(edgetpu_path, "test_data/mobilenet_v1_1.0_224_quant.tflite")
+ return model_path
+
+ def init_interpreter(model_path, target_edgetpu):
+ # Initialize interpreter
+ if target_edgetpu:
+ edgetpu_path = os.getenv('EDGETPU_PATH', "/home/mendel/edgetpu")
+ libedgetpu = os.path.join(edgetpu_path, "libedgetpu/direct/aarch64/libedgetpu.so.1")
+ interpreter = tflite.Interpreter(
+ model_path=model_path,
+ experimental_delegates=[tflite.load_delegate(libedgetpu)])
+ else:
+ interpreter = tflite.Interpreter(model_path=model_path)
+ return interpreter
+
+ def check_remote(target_edgetpu=False):
+ tflite_model_path = get_tflite_model_path(target_edgetpu)
+
+ # inference via tflite interpreter python apis
+ interpreter = init_interpreter(tflite_model_path, target_edgetpu)
+ interpreter.allocate_tensors()
+ input_details = interpreter.get_input_details()
+ output_details = interpreter.get_output_details()
+
+ input_shape = input_details[0]['shape']
+ tflite_input = np.array(np.random.random_sample(input_shape), dtype=np.uint8)
+ interpreter.set_tensor(input_details[0]['index'], tflite_input)
+ interpreter.invoke()
+ tflite_output = interpreter.get_tensor(output_details[0]['index'])
+
+ # inference via remote tvm tflite runtime
+ server = rpc.Server("localhost")
+ remote = rpc.connect(server.host, server.port)
+ ctx = remote.cpu(0)
+
+ with open(tflite_model_path, 'rb') as model_fin:
+ runtime = tflite_runtime.create(model_fin.read(), ctx)
+ runtime.set_input(0, tvm.nd.array(tflite_input, ctx))
+ runtime.invoke()
+ out = runtime.get_output(0)
+ np.testing.assert_equal(out.asnumpy(), tflite_output)
+
+ # Target CPU on coral board
+ check_remote()
+ # Target EdgeTPU on coral board
+ check_remote(target_edgetpu=True)
+
+if __name__ == "__main__":
+ # skipped_test_tflite_runtime()
+ pass
return tflite_model
- def check_verify():
+ def check_local():
tflite_fname = "model.tflite"
tflite_model = create_tflite_model()
temp = util.tempdir()
tflite_model_path = temp.relpath(tflite_fname)
- print(tflite_model_path)
open(tflite_model_path, 'wb').write(tflite_model)
# inference via tflite interpreter python apis
- print('interpreter')
interpreter = tflite.Interpreter(model_path=tflite_model_path)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
interpreter.invoke()
tflite_output = interpreter.get_tensor(output_details[0]['index'])
- print('tvm tflite runtime')
# inference via tvm tflite runtime
with open(tflite_model_path, 'rb') as model_fin:
runtime = tflite_runtime.create(model_fin.read(), tvm.cpu(0))
- runtime.allocate_tensors()
runtime.set_input(0, tvm.nd.array(tflite_input))
runtime.invoke()
out = runtime.get_output(0)
with open(tflite_model_path, 'rb') as model_fin:
runtime = tflite_runtime.create(model_fin.read(), remote.cpu(0))
- runtime.allocate_tensors()
runtime.set_input(0, tvm.nd.array(tflite_input, remote.cpu(0)))
runtime.invoke()
out = runtime.get_output(0)
np.testing.assert_equal(out.asnumpy(), tflite_output)
-
- check_verify()
+ check_local()
check_remote()
if __name__ == "__main__":
projects / platform / upstream / tvm.git / commitdiff