From 0ea54b14615093e7fcb7c2cd441f828d102f161c Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Thu, 29 Jun 2023 14:11:29 -0700 Subject: [PATCH] Add delegate options to base options for java API. and add unit tset for BaseOptions. PiperOrigin-RevId: 544458644 --- mediapipe/tasks/cc/core/base_options.cc | 12 +- mediapipe/tasks/cc/core/base_options_test.cc | 5 +- .../mediapipe/tasks/core/BaseOptions.java | 81 +++++++++ .../mediapipe/tasks/core/TaskOptions.java | 34 ++++ .../mediapipe/tasks/core/AndroidManifest.xml | 24 +++ .../com/google/mediapipe/tasks/core/BUILD | 2 + .../mediapipe/tasks/core/BaseOptionsTest.java | 159 ++++++++++++++++++ 7 files changed, 312 insertions(+), 5 deletions(-) create mode 100644 mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/AndroidManifest.xml create mode 100644 mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/BaseOptionsTest.java diff --git a/mediapipe/tasks/cc/core/base_options.cc b/mediapipe/tasks/cc/core/base_options.cc index b7987f98..863338fe 100644 --- a/mediapipe/tasks/cc/core/base_options.cc +++ b/mediapipe/tasks/cc/core/base_options.cc @@ -41,9 +41,15 @@ proto::Acceleration ConvertDelegateOptionsToAccelerationProto( proto::Acceleration acceleration_proto = proto::Acceleration(); auto* gpu = acceleration_proto.mutable_gpu(); gpu->set_use_advanced_gpu_api(true); - gpu->set_cached_kernel_path(options.cached_kernel_path); - gpu->set_serialized_model_dir(options.serialized_model_dir); - gpu->set_model_token(options.model_token); + if (!options.cached_kernel_path.empty()) { + gpu->set_cached_kernel_path(options.cached_kernel_path); + } + if (!options.serialized_model_dir.empty()) { + gpu->set_serialized_model_dir(options.serialized_model_dir); + } + if (!options.model_token.empty()) { + gpu->set_model_token(options.model_token); + } return acceleration_proto; } diff --git a/mediapipe/tasks/cc/core/base_options_test.cc b/mediapipe/tasks/cc/core/base_options_test.cc index af9a55a3..39066351 100644 --- a/mediapipe/tasks/cc/core/base_options_test.cc +++ b/mediapipe/tasks/cc/core/base_options_test.cc @@ -59,14 +59,15 @@ TEST(DelegateOptionsTest, SucceedGpuOptions) { BaseOptions base_options; base_options.delegate = BaseOptions::Delegate::GPU; BaseOptions::GpuOptions gpu_options; - gpu_options.cached_kernel_path = kCachedModelDir; + gpu_options.serialized_model_dir = kCachedModelDir; gpu_options.model_token = kModelToken; base_options.delegate_options = gpu_options; proto::BaseOptions proto = ConvertBaseOptionsToProto(&base_options); ASSERT_TRUE(proto.acceleration().has_gpu()); ASSERT_FALSE(proto.acceleration().has_tflite()); EXPECT_TRUE(proto.acceleration().gpu().use_advanced_gpu_api()); - EXPECT_EQ(proto.acceleration().gpu().cached_kernel_path(), kCachedModelDir); + EXPECT_FALSE(proto.acceleration().gpu().has_cached_kernel_path()); + EXPECT_EQ(proto.acceleration().gpu().serialized_model_dir(), kCachedModelDir); EXPECT_EQ(proto.acceleration().gpu().model_token(), kModelToken); } diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/BaseOptions.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/BaseOptions.java index 8eec72ef..dc2c001b 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/BaseOptions.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/BaseOptions.java @@ -54,6 +54,9 @@ public abstract class BaseOptions { */ public abstract Builder setDelegate(Delegate delegate); + /** Options for the chosen delegate. If not set, the default delegate options is used. */ + public abstract Builder setDelegateOptions(DelegateOptions delegateOptions); + abstract BaseOptions autoBuild(); /** @@ -79,6 +82,23 @@ public abstract class BaseOptions { throw new IllegalArgumentException( "The model buffer should be either a direct ByteBuffer or a MappedByteBuffer."); } + boolean delegateMatchesDelegateOptions = true; + if (options.delegateOptions().isPresent()) { + switch (options.delegate()) { + case CPU: + delegateMatchesDelegateOptions = + options.delegateOptions().get() instanceof DelegateOptions.CpuOptions; + break; + case GPU: + delegateMatchesDelegateOptions = + options.delegateOptions().get() instanceof DelegateOptions.GpuOptions; + break; + } + if (!delegateMatchesDelegateOptions) { + throw new IllegalArgumentException( + "Specified Delegate type does not match the provided delegate options."); + } + } return options; } } @@ -91,6 +111,67 @@ public abstract class BaseOptions { abstract Delegate delegate(); + abstract Optional delegateOptions(); + + /** Advanced config options for the used delegate. */ + public abstract static class DelegateOptions { + + /** Options for CPU. */ + @AutoValue + public abstract static class CpuOptions extends DelegateOptions { + + public static Builder builder() { + Builder builder = new AutoValue_BaseOptions_DelegateOptions_CpuOptions.Builder(); + return builder; + } + + /** Builder for {@link CpuOptions}. */ + @AutoValue.Builder + public abstract static class Builder { + + public abstract CpuOptions build(); + } + } + + /** Options for GPU. */ + @AutoValue + public abstract static class GpuOptions extends DelegateOptions { + // Load pre-compiled serialized binary cache to accelerate init process. + // Only available on Android. Kernel caching will only be enabled if this + // path is set. NOTE: binary cache usage may be skipped if valid serialized + // model, specified by "serialized_model_dir", exists. + abstract Optional cachedKernelPath(); + + // A dir to load from and save to a pre-compiled serialized model used to + // accelerate init process. + // NOTE: serialized model takes precedence over binary cache + // specified by "cached_kernel_path", which still can be used if + // serialized model is invalid or missing. + abstract Optional serializedModelDir(); + + // Unique token identifying the model. Used in conjunction with + // "serialized_model_dir". It is the caller's responsibility to ensure + // there is no clash of the tokens. + abstract Optional modelToken(); + + public static Builder builder() { + return new AutoValue_BaseOptions_DelegateOptions_GpuOptions.Builder(); + } + + /** Builder for {@link GpuOptions}. */ + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setCachedKernelPath(String cachedKernelPath); + + public abstract Builder setSerializedModelDir(String serializedModelDir); + + public abstract Builder setModelToken(String modelToken); + + public abstract GpuOptions build(); + } + } + } + public static Builder builder() { return new AutoValue_BaseOptions.Builder().setDelegate(Delegate.CPU); } diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/TaskOptions.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/TaskOptions.java index 11330ac0..991aceba 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/core/TaskOptions.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/core/TaskOptions.java @@ -61,17 +61,51 @@ public abstract class TaskOptions { accelerationBuilder.setTflite( InferenceCalculatorProto.InferenceCalculatorOptions.Delegate.TfLite .getDefaultInstance()); + options + .delegateOptions() + .ifPresent( + delegateOptions -> + setDelegateOptions( + accelerationBuilder, + (BaseOptions.DelegateOptions.CpuOptions) delegateOptions)); break; case GPU: accelerationBuilder.setGpu( InferenceCalculatorProto.InferenceCalculatorOptions.Delegate.Gpu.newBuilder() .setUseAdvancedGpuApi(true) .build()); + options + .delegateOptions() + .ifPresent( + delegateOptions -> + setDelegateOptions( + accelerationBuilder, + (BaseOptions.DelegateOptions.GpuOptions) delegateOptions)); break; } + return BaseOptionsProto.BaseOptions.newBuilder() .setModelAsset(externalFileBuilder.build()) .setAcceleration(accelerationBuilder.build()) .build(); } + + private void setDelegateOptions( + AccelerationProto.Acceleration.Builder accelerationBuilder, + BaseOptions.DelegateOptions.CpuOptions options) { + accelerationBuilder.setTflite( + InferenceCalculatorProto.InferenceCalculatorOptions.Delegate.TfLite.getDefaultInstance()); + } + + private void setDelegateOptions( + AccelerationProto.Acceleration.Builder accelerationBuilder, + BaseOptions.DelegateOptions.GpuOptions options) { + InferenceCalculatorProto.InferenceCalculatorOptions.Delegate.Gpu.Builder gpuBuilder = + InferenceCalculatorProto.InferenceCalculatorOptions.Delegate.Gpu.newBuilder() + .setUseAdvancedGpuApi(true); + options.cachedKernelPath().ifPresent(gpuBuilder::setCachedKernelPath); + options.serializedModelDir().ifPresent(gpuBuilder::setSerializedModelDir); + options.modelToken().ifPresent(gpuBuilder::setModelToken); + accelerationBuilder.setGpu(gpuBuilder.build()); + } } diff --git a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/AndroidManifest.xml b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/AndroidManifest.xml new file mode 100644 index 00000000..26310fc1 --- /dev/null +++ b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/AndroidManifest.xml @@ -0,0 +1,24 @@ + + + + + + + + + + + + + + + diff --git a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/BUILD b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/BUILD index 01e7ad0f..ce7435d6 100644 --- a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/BUILD +++ b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/BUILD @@ -23,3 +23,5 @@ android_library( "//third_party/java/android_libs/guava_jdk5:io", ], ) + +# TODO: Enable this in OSS diff --git a/mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/BaseOptionsTest.java b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/BaseOptionsTest.java new file mode 100644 index 00000000..939ecb40 --- /dev/null +++ b/mediapipe/tasks/javatests/com/google/mediapipe/tasks/core/BaseOptionsTest.java @@ -0,0 +1,159 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed 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. + +package com.google.mediapipe.tasks.core; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import androidx.test.ext.junit.runners.AndroidJUnit4; +import com.google.mediapipe.proto.CalculatorOptionsProto.CalculatorOptions; +import com.google.mediapipe.tasks.core.proto.AccelerationProto; +import com.google.mediapipe.tasks.core.proto.BaseOptionsProto; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Suite; +import org.junit.runners.Suite.SuiteClasses; + +/** Test for {@link BaseOptions} */ +@RunWith(Suite.class) +@SuiteClasses({BaseOptionsTest.General.class, BaseOptionsTest.ConvertProtoTest.class}) +public class BaseOptionsTest { + + static final String MODEL_ASSET_PATH = "dummy_model.tflite"; + static final String SERIALIZED_MODEL_DIR = "dummy_serialized_model_dir"; + static final String MODEL_TOKEN = "dummy_model_token"; + static final String CACHED_KERNEL_PATH = "dummy_cached_kernel_path"; + + @RunWith(AndroidJUnit4.class) + public static final class General extends BaseOptionsTest { + @Test + public void succeedsWithDefaultOptions() throws Exception { + BaseOptions options = BaseOptions.builder().setModelAssetPath(MODEL_ASSET_PATH).build(); + assertThat(options.modelAssetPath().isPresent()).isTrue(); + assertThat(options.modelAssetPath().get()).isEqualTo(MODEL_ASSET_PATH); + assertThat(options.delegate()).isEqualTo(Delegate.CPU); + } + + @Test + public void succeedsWithGpuOptions() throws Exception { + BaseOptions options = + BaseOptions.builder() + .setModelAssetPath(MODEL_ASSET_PATH) + .setDelegate(Delegate.GPU) + .setDelegateOptions( + BaseOptions.DelegateOptions.GpuOptions.builder() + .setSerializedModelDir(SERIALIZED_MODEL_DIR) + .setModelToken(MODEL_TOKEN) + .setCachedKernelPath(CACHED_KERNEL_PATH) + .build()) + .build(); + assertThat( + ((BaseOptions.DelegateOptions.GpuOptions) options.delegateOptions().get()) + .serializedModelDir() + .get()) + .isEqualTo(SERIALIZED_MODEL_DIR); + assertThat( + ((BaseOptions.DelegateOptions.GpuOptions) options.delegateOptions().get()) + .modelToken() + .get()) + .isEqualTo(MODEL_TOKEN); + assertThat( + ((BaseOptions.DelegateOptions.GpuOptions) options.delegateOptions().get()) + .cachedKernelPath() + .get()) + .isEqualTo(CACHED_KERNEL_PATH); + } + + @Test + public void failsWithInvalidDelegateOptions() throws Exception { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + BaseOptions.builder() + .setModelAssetPath(MODEL_ASSET_PATH) + .setDelegate(Delegate.CPU) + .setDelegateOptions( + BaseOptions.DelegateOptions.GpuOptions.builder() + .setSerializedModelDir(SERIALIZED_MODEL_DIR) + .setModelToken(MODEL_TOKEN) + .build()) + .build()); + assertThat(exception) + .hasMessageThat() + .contains("Specified Delegate type does not match the provided delegate options."); + } + } + + /** A mock TaskOptions class providing access to convertBaseOptionsToProto. */ + public static class MockTaskOptions extends TaskOptions { + + public MockTaskOptions(BaseOptions baseOptions) { + baseOptionsProto = convertBaseOptionsToProto(baseOptions); + } + + public BaseOptionsProto.BaseOptions getBaseOptionsProto() { + return baseOptionsProto; + } + + private BaseOptionsProto.BaseOptions baseOptionsProto; + + @Override + public CalculatorOptions convertToCalculatorOptionsProto() { + return CalculatorOptions.newBuilder().build(); + } + } + + /** Test for converting {@link BaseOptions} to {@link BaseOptionsProto} */ + @RunWith(AndroidJUnit4.class) + public static final class ConvertProtoTest extends BaseOptionsTest { + @Test + public void succeedsWithDefaultOptions() throws Exception { + BaseOptions options = + BaseOptions.builder() + .setModelAssetPath(MODEL_ASSET_PATH) + .setDelegate(Delegate.CPU) + .setDelegateOptions(BaseOptions.DelegateOptions.CpuOptions.builder().build()) + .build(); + MockTaskOptions taskOptions = new MockTaskOptions(options); + AccelerationProto.Acceleration acceleration = + taskOptions.getBaseOptionsProto().getAcceleration(); + assertThat(acceleration.hasTflite()).isTrue(); + } + + @Test + public void succeedsWithGpuOptions() throws Exception { + BaseOptions options = + BaseOptions.builder() + .setModelAssetPath(MODEL_ASSET_PATH) + .setDelegate(Delegate.GPU) + .setDelegateOptions( + BaseOptions.DelegateOptions.GpuOptions.builder() + .setModelToken(MODEL_TOKEN) + .setSerializedModelDir(SERIALIZED_MODEL_DIR) + .build()) + .build(); + MockTaskOptions taskOptions = new MockTaskOptions(options); + AccelerationProto.Acceleration acceleration = + taskOptions.getBaseOptionsProto().getAcceleration(); + assertThat(acceleration.hasTflite()).isFalse(); + assertThat(acceleration.hasGpu()).isTrue(); + assertThat(acceleration.getGpu().getUseAdvancedGpuApi()).isTrue(); + assertThat(acceleration.getGpu().hasCachedKernelPath()).isFalse(); + assertThat(acceleration.getGpu().getModelToken()).isEqualTo(MODEL_TOKEN); + assertThat(acceleration.getGpu().getSerializedModelDir()).isEqualTo(SERIALIZED_MODEL_DIR); + } + } +}