pytorch-lite-multiplatform

Kotlin 在 Android 和 iOS 上的 PyTorch Lite 库的多平台包装器。您可以使用此库在您的 Kotlin 多平台项目中编写 PyTorch Lite 模型的移动推理代码。API 与 PyTorch Lite 的 Android API 非常接近。

安装

将以下内容添加到您的 shared/build.gradle.kts 作为 commonMain 依赖项。

implementation("de.voize:pytorch-lite-multiplatform:<version>")

在 shared/build.gradle.kts 中的 cocoapods 插件块中添加 PLMLibTorchWrapper pod，并添加 useLibraries()，因为 PLMLibTorchWrapper pod 依赖于包含静态库的 LibTorch-Lite pod。

cocoapods {
    ...

    pod("PLMLibTorchWrapper") {
        version = "<version>"
        headers = "LibTorchWrapper.h"
    }

    useLibraries()
}

如果您使用的是 Kotlin 版本 < 1.8.0，则 headers 属性不可用。相反，您必须将以下内容添加到您的 shared/build.gradle.kts 中（有关更多信息，请参阅 此问题）

tasks.named<org.jetbrains.kotlin.gradle.tasks.DefFileTask>("generateDefPLMLibTorchWrapper").configure {
    doLast {
        outputFile.writeText("""
            language = Objective-C
            headers = LibTorchWrapper.h
        """.trimIndent())
    }
}

其他步骤

• 确保在您的 iOS XCode 项目中禁用了位码
• 确保您的 iOS 应用 Podfile 中不包含 use_frameworks!
• 您的 framework 块可能需要声明 isStatic = true

使用说明

首先，将 PyTorch 模型导出为轻量级解释器格式。管理您的应用程序中导出模型文件的存储，例如与您的应用程序捆绑、在应用程序初始化期间从服务器下载，或其他方式。然后，您可以使用模型文件的路径来初始化 TorchModule。

import de.voize.pytorch_lite_multiplatform.TorchModule

val module = TorchModule(path = "<path/to/model.ptl>")

一旦初始化了模型，您就可以开始进行推理了。

类似于 PyTorch Lite 的 Android API，您可以使用 IValue 和 Tensor 将输入数据传递给您的模型，并处理模型输出。为了管理为您 tensor 分配的内存，您需要使用 plmScoped 来指定内存分配需要保持的最远点。

import de.voize.pytorch_lite_multiplatform.*

plmScoped {
    val inputTensor = Tensor.fromBlob(
        data = floatArrayOf(...),
        shape = longArrayOf(...),
        scope = this
    )

    val inputIValue = IValue.fromTensor(inputTensor)

    val output = module.forward(inputIValue)
    // you could also use
    // module.runMethod("forward", inputIValue)

    val outputTensor = output.toTensor()
    val outputData = outputTensor.getDataAsFloatArray()

    ...
}

IValue在构建所需模型输入（例如张量、标量、标志、字典、元组等）方面非常灵活。有关所有可用选项，请参阅IValue接口，并通过PyTorch 的 Android 示例应用程序查找使用 IValue 进行推理的示例。

内存管理

为了简化在 Android 和 iOS 之间分配推理资源的管理，我们引入了PLMScope和plmScoped实用工具。在 Android 上，JVM垃圾回收和PyTorch Lite很好地管理分配的内存，因此plmScoped是一个空操作。但在 iOS 上，内存是在 Kotlin 中分配的，并与原生Objective-C代码进行交换，反之亦然，没有自动释放资源。这就是plmScoped发挥作用并释放推理分配的内存的地方。因此，您需要正确定义资源需要保持分配作用的范围，以免出现内存泄露或丢失未来需要的内存。

运行测试

iOS

要在iOS上运行测试，执行 iosSimulatorX64Test gradle 任务

./gradlew iosSimulatorX64Test

这将自动调用 build_dummy_model.py 来创建用于测试的模拟 torchscript 模块，将其复制到模拟文件目录并执行测试。请确保选择一个其中包含 torch 依赖的 Python 环境。