TorchMobileでGPUパワーを使う

MobileNetv2をiPhone11で実行。
CPUのみ、CoreMLバックエンド、Metalバックエンドを比べてみました。

PyTorchモバイルを高速で使いたい

CoreMLバックエンドもしくはMetalが使える

CoreMLバックエンドやMetalを使用するようにモデルを設定できます。
これでデバイスのGPUやニューラルエンジンを使えるようです。

変換方法

CPU

import torch
import torchvision
from torch.utils.mobile_optimizer import optimize_for_mobile

model = torchvision.models.mobilenet_v2(pretrained=True)
model.eval()
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
torchscript_model_optimized = optimize_for_mobile(traced_script_module)
torchscript_model_optimized._save_for_lite_interpreter("HelloWorld/HelloWorld/model/model.pt")

CoreML

mobilenetv2.py

import torch
import torchvision

from torch.backends._coreml.preprocess import (
    CompileSpec,
    TensorSpec,
    CoreMLComputeUnit,
)

def mobilenetv2_spec():
    return {
        "forward": CompileSpec(
            inputs=(
                TensorSpec(
                    shape=[1, 3, 224, 224],
                ),
            ),
            outputs=(
                TensorSpec(
                    shape=[1, 1000],
                ),
            ),
            backend=CoreMLComputeUnit.ALL,
            allow_low_precision=True,
        ),
    }


def main():
    model = torchvision.models.mobilenet_v2(pretrained=True)
    model.eval()
    example = torch.rand(1, 3, 224, 224)
    model = torch.jit.trace(model, example)
    compile_spec = mobilenetv2_spec()
    mlmodel = torch._C._jit_to_backend("coreml", model, compile_spec)
    mlmodel._save_for_lite_interpreter("./mobilenetv2_coreml.ptl")


if __name__ == "__main__":
    main()

Metal

import torch
import torchvision
from torch.utils.mobile_optimizer import optimize_for_mobile

model = torchvision.models.mobilenet_v2(pretrained=True)
scripted_model = torch.jit.script(model)
optimized_model = optimize_for_mobile(scripted_model, backend='metal')
print(torch.jit.export_opnames(optimized_model))
optimized_model._save_for_lite_interpreter('./mobilenetv2_metal.pt')

推論

モデル実行の手順は同じです。

モデルインスタンス化

try {
  _impl = torch::jit::_load_for_mobile(filePath.UTF8String);
} catch (const std::exception& exception) {
  NSLog(@"%s", exception.what());
  return nil;
}

実行

auto outputTensor = _impl.forward({tensor}).toTensor().cpu();

比較結果

MobileNetv2をiPhone11で実行。

Model	Time
CPU	0.024118
CoreML	0.013819
Metal	0.046400

CoreMLが倍速かったです。
Metalバックエンドがそんなに速くないのは意外でした。

🐣

フリーランスエンジニアです。
お仕事のご相談こちらまで
rockyshikoku@gmail.com

Core MLやARKitを使ったアプリを作っています。
機械学習／AR関連の情報を発信しています。

Twitter
Medium