Android端PyTorch模型推理加速实战：从模型优化到硬件加速

1次阅读

共计 2179 个字符，预计需要花费 6 分钟才能阅读完成。

在 Android 设备上部署 PyTorch 模型时，开发者通常会遇到以下核心问题：

计算能力限制：移动端 CPU 的算力仅为服务器的 1 /10~1/100，ResNet50 等常见模型单次推理可能达到 300-500ms
内存瓶颈：模型参数和中间激活值占用大量内存，512MB 以下设备易出现 OOM
功耗敏感：持续高负载运算导致发热降频，实测显示 CPU 满负载时推理延迟可能增加 2 - 3 倍

对比主流推理框架在 Android 端的表现：

框架	模型兼容性	量化支持	硬件加速	内存占用
TensorFlow Lite	★★★★☆	★★★★★	★★★★☆	★★★☆☆
PyTorch Mobile	★★★★★	★★★★☆	★★★☆☆	★★★★☆
ONNX Runtime	★★★☆☆	★★★★☆	★★★★☆	★★★☆☆

PyTorch Mobile 的核心优势在于：
– 原生支持 TorchScript 模型，无需额外转换步骤
– 与训练代码无缝衔接，支持动态图特性
– 2022 年后显著优化了算子覆盖率和内存管理

# 原始模型导出
model = resnet18(pretrained=True)
model.eval()
traced_script_module = torch.jit.trace(model, torch.rand(1,3,224,224))

# 动态量化实施
quantized_model = torch.quantization.quantize_dynamic(
    traced_script_module,
    {torch.nn.Linear},  # 量化目标层类型
    dtype=torch.qint8
)
quantized_model.save("quantized_resnet18.pt")

插入量化 / 反量化节点（torch.quantization.QuantStub）
融合 Conv+ReLU 等算子组合
校准模型（500-1000 张校准数据）
转换为量化模型

Android 端需添加依赖：

dependencies {
    implementation 'org.pytorch:pytorch_android:1.12.0'
    implementation 'org.pytorch:pytorch_android_torchvision:1.12.0'
    implementation 'org.pytorch:pytorch_android_gpu:1.12.0'  // GPU 支持
}

推理代码适配：

val module = Module.load(assetFilePath(this, "model.pt"), Device.GPU)
val inputTensor = TensorImageUtils.bitmapToFloat32Tensor(
    bitmap,
    floatArrayOf(0.485f, 0.456f, 0.406f), 
    floatArrayOf(0.229f, 0.224f, 0.225f)
).to(Device.GPU)

val nnapiModule = NnapiModule.load(context.assets.open("model.nnc"),  // 需先转换模型
    Device.CPU
)
// 输入数据需满足 NCHW 内存布局
val inputBuffer = ByteBuffer.allocateDirect(1*3*224*224*4).order(ByteOrder.nativeOrder())

测试设备：Pixel 6 (Tensor G1)
| 方案 | 延迟(ms) | 内存(MB) | 功耗(mW) |
|—————|———-|———-|———-|
| FP32 CPU | 342 | 287 | 2100 |
| INT8 CPU | 89 | 152 | 950 |
| FP16 GPU | 64 | 203 | 1800 |
| NNAPI (INT8) | 47 | 118 | 620 |

报错：”Unsupported operator aten::upsample_bilinear2d”
解决：使用torch._C._jit_pass_onnx_scalar_type_analysis 进行算子兼容性检查

class InferenceSession : AutoCloseable {
    private val nativeHandle: Long

    init {nativeHandle = initNative()
        MemoryMonitor.register(this)
    }

    override fun close() {releaseNative(nativeHandle)
    }
}