experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add single-parameter versions of AddBuiltin and AddCustom functions.

Browse Source 
This allows us to make the max_version the same as min_version by default; C++ does not allow the default value of a parameter to depend on the other parameters.

PiperOrigin-RevId: 301269532
Change-Id: If2023ec1a2f7081e601bb95fd65f0f52c6eb83c6
This commit is contained in:
Robert David 2020-03-16 17:16:38 -07:00 committed by TensorFlower Gardener
parent eb7bdc25ca
commit 876a3c1708
5 changed files with 86 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
tensorflow/lite/micro/micro_mutable_op_resolver.h
View File

@ -61,37 +61,45 @@ class MicroOpResolver : public OpResolver {
}
void AddBuiltin(tflite::BuiltinOperator op, TfLiteRegistration* registration,
int min_version = 1, int max_version = 1) {
for (int version = min_version; version <= max_version; ++version) {
if (registrations_len_ >= tOpCount) {
// TODO(b/147748244) - Add error reporting hooks so we can report this!
return;
}
TfLiteRegistration* new_registration =
&registrations_[registrations_len_];
registrations_len_ += 1;
int version = 1) {
if (registrations_len_ >= tOpCount) {
// TODO(b/147748244) - Add error reporting hooks so we can report this!
return;
}
TfLiteRegistration* new_registration = &registrations_[registrations_len_];
registrations_len_ += 1;
*new_registration = *registration;
new_registration->builtin_code = op;
new_registration->version = version;
*new_registration = *registration;
new_registration->builtin_code = op;
new_registration->version = version;
}
void AddBuiltin(tflite::BuiltinOperator op, TfLiteRegistration* registration,
int min_version, int max_version) {
for (int version = min_version; version <= max_version; ++version) {
AddBuiltin(op, registration, version);
}
}
void AddCustom(const char* name, TfLiteRegistration* registration,
int min_version = 1, int max_version = 1) {
for (int version = min_version; version <= max_version; ++version) {
if (registrations_len_ >= tOpCount) {
// TODO(b/147748244) - Add error reporting hooks so we can report this!
return;
}
TfLiteRegistration* new_registration =
&registrations_[registrations_len_];
registrations_len_ += 1;
int version = 1) {
if (registrations_len_ >= tOpCount) {
// TODO(b/147748244) - Add error reporting hooks so we can report this!
return;
}
TfLiteRegistration* new_registration = &registrations_[registrations_len_];
registrations_len_ += 1;
*new_registration = *registration;
new_registration->builtin_code = BuiltinOperator_CUSTOM;
new_registration->custom_name = name;
new_registration->version = version;
*new_registration = *registration;
new_registration->builtin_code = BuiltinOperator_CUSTOM;
new_registration->custom_name = name;
new_registration->version = version;
}
void AddCustom(const char* name, TfLiteRegistration* registration,
int min_version, int max_version) {
for (int version = min_version; version <= max_version; ++version) {
AddCustom(name, registration, version);
}
}

3
tensorflow/lite/micro/micro_mutable_op_resolver_test.cc
View File

@ -113,6 +113,9 @@ TF_LITE_MICRO_TEST(TestZeroVersionRegistration) {
MicroOpResolver<1> micro_op_resolver;
micro_op_resolver.AddCustom("mock_custom", &r,
tflite::MicroOpResolverAnyVersion());
TF_LITE_MICRO_EXPECT_EQ(1, micro_op_resolver.GetRegistrationLength());
OpResolver* resolver = &micro_op_resolver;
const TfLiteRegistration* registration = resolver->FindOp("mock_custom", 0);

36
tensorflow/lite/mutable_op_resolver.cc
View File

@ -29,29 +29,41 @@ const TfLiteRegistration* MutableOpResolver::FindOp(const char* op,
return it != custom_ops_.end() ? &it->second : nullptr;
}
void MutableOpResolver::AddBuiltin(tflite::BuiltinOperator op,
const TfLiteRegistration* registration,
int version) {
TfLiteRegistration new_registration = *registration;
new_registration.custom_name = nullptr;
new_registration.builtin_code = op;
new_registration.version = version;
auto op_key = std::make_pair(op, version);
builtins_[op_key] = new_registration;
}
void MutableOpResolver::AddBuiltin(tflite::BuiltinOperator op,
const TfLiteRegistration* registration,
int min_version, int max_version) {
for (int version = min_version; version <= max_version; ++version) {
TfLiteRegistration new_registration = *registration;
new_registration.custom_name = nullptr;
new_registration.builtin_code = op;
new_registration.version = version;
auto op_key = std::make_pair(op, version);
builtins_[op_key] = new_registration;
AddBuiltin(op, registration, version);
}
}
void MutableOpResolver::AddCustom(const char* name,
const TfLiteRegistration* registration,
int version) {
TfLiteRegistration new_registration = *registration;
new_registration.builtin_code = BuiltinOperator_CUSTOM;
new_registration.custom_name = name;
new_registration.version = version;
auto op_key = std::make_pair(name, version);
custom_ops_[op_key] = new_registration;
}
void MutableOpResolver::AddCustom(const char* name,
const TfLiteRegistration* registration,
int min_version, int max_version) {
for (int version = min_version; version <= max_version; ++version) {
TfLiteRegistration new_registration = *registration;
new_registration.builtin_code = BuiltinOperator_CUSTOM;
new_registration.custom_name = name;
new_registration.version = version;
auto op_key = std::make_pair(name, version);
custom_ops_[op_key] = new_registration;
AddCustom(name, registration, version);
}
}

10
tensorflow/lite/mutable_op_resolver.h
View File

@ -60,10 +60,14 @@ class MutableOpResolver : public OpResolver {
int version) const override;
const TfLiteRegistration* FindOp(const char* op, int version) const override;
void AddBuiltin(tflite::BuiltinOperator op,
const TfLiteRegistration* registration, int min_version = 1,
int max_version = 1);
const TfLiteRegistration* registration, int version = 1);
void AddBuiltin(tflite::BuiltinOperator op,
const TfLiteRegistration* registration, int min_version,
int max_version);
void AddCustom(const char* name, const TfLiteRegistration* registration,
int min_version = 1, int max_version = 1);
int version = 1);
void AddCustom(const char* name, const TfLiteRegistration* registration,
int min_version, int max_version);
void AddAll(const MutableOpResolver& other);
private:

19
tensorflow/lite/mutable_op_resolver_test.cc
View File

@ -81,6 +81,25 @@ TEST(MutableOpResolverTest, FindMissingOp) {
EXPECT_EQ(found_registration, nullptr);
}
TEST(MutableOpResolverTest, RegisterOpWithSingleVersion) {
MutableOpResolver resolver;
// The kernel supports version 2 only
resolver.AddBuiltin(BuiltinOperator_ADD, GetDummyRegistration(), 2);
const TfLiteRegistration* found_registration;
found_registration = resolver.FindOp(BuiltinOperator_ADD, 1);
ASSERT_EQ(found_registration, nullptr);
found_registration = resolver.FindOp(BuiltinOperator_ADD, 2);
ASSERT_NE(found_registration, nullptr);
EXPECT_TRUE(found_registration->invoke == DummyInvoke);
EXPECT_EQ(found_registration->version, 2);
found_registration = resolver.FindOp(BuiltinOperator_ADD, 3);
ASSERT_EQ(found_registration, nullptr);
}
TEST(MutableOpResolverTest, RegisterOpWithMultipleVersions) {
MutableOpResolver resolver;
// The kernel supports version 2 and 3