TensorFlow LiteのGPUのソースコードが公開されたよ

Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

このブログで1月21日に書いたTensorFlow LiteでGPUが使えるようになったのソースコードが公開されました。

TFLite on GPU

「iOSのMetal」と「AndroidのOpenGL v3」用のコードが公開されました。

DeleGateにて、上手く切り替えています。

////////
// Set up interpreter.
auto model = FlatBufferModel::BuildFromFile(model_path);
ops::builtin::BuiltinOpResolver op_resolver;
std::unique_ptr<Interpreter> interpreter;
InterpreterBuilder(*model, op_resolver)(&interpreter);

////////
// NEW: Prepare GPU delegate.
auto* delegate = TfLiteGpuDelegateCreate(/*options=*/nullptr);
if (interpreter->ModifyGraphWithDelegate(delegate) != kTfLiteOk) return;

////////
// Run inference.
WriteToInputTensor(interpreter->typed_input_tensor<float>(0));
if (interpreter->Invoke() != kTfLiteOk) return;
ReadFromOutputTensor(interpreter->typed_output_tensor<float>(0));

////////
// Clean up.
TfLiteGpuDelegateDelete(delegate);

　Androidでは、=> TFLite GPU backend uses OpenGL compute shaders and thus requires OpenGL ES 3.1 or higher.

　iOSでは、=> Metal shaders are used for iOS, which were introduced with iOS 8.

のようです。

BUILDを見ると、

Androidでは、

# build -c opt --config android_arm64 --copt -Os --copt -DTFLITE_GPU_BINARY_RELEASE --copt -fvisibility=hidden --linkopt -s --strip always :libtflite_gpu_gl.so
cc_binary(
    name = "libtflite_gpu_gl.so",
    linkopts = select({
        "//tensorflow:android": [
            "-lEGL",
            "-lGLESv3",
        ],
        "//conditions:default": [],
    }),
    linkshared = 1,
    linkstatic = 1,
    tags = [
        "nobuilder",
        "notap",
    ],
    deps = [":gl_delegate"],
)

のように、ARM64で、OpenGL ES v3 です。

一方、iOS では、64ビットです。

# build -c opt --config ios_arm64 --copt -Os --copt -DTFLITE_GPU_BINARY_RELEASE --copt -fvisibility=hidden --linkopt -s --strip always :libtflite_gpu_metal.so
cc_binary(
    name = "libtflite_gpu_metal.so",
    linkshared = 1,
    linkstatic = 1,
    tags = [
        "nobuilder",
        "notap",
    ],
    deps = [":metal_delegate"],
)