Ubuntu PyTorch如何进行模型部署到移动设备

在Ubuntu上使用PyTorch将模型部署到移动设备，主要流程包括模型转换、优化及适配不同平台，具体步骤如下：

一、模型转换

1. 导出为TorchScript格式
   使用torch.jit.trace或torch.jit.script将训练好的模型转换为TorchScript格式，适配PyTorch Mobile。

   import torch
   from your_model import YourModel
   
   model = YourModel().eval()
   example_input = torch.rand(1, 3, 224, 224)  # 根据模型输入调整
   traced_model = torch.jit.trace(model, example_input)
   traced_model.save("model.pt")  # 保存为.pt文件
   

2. 可选：转换为ONNX格式
   若需跨平台或进一步优化，可导出为ONNX格式，但需注意PyTorch Mobile对ONNX的支持有限，建议优先使用TorchScript。

   torch.onnx.export(model, example_input, "model.onnx", input_names=["input"], output_names=["output"])
   

二、模型优化

1. 量化压缩
   通过动态量化或静态量化减少模型体积、提升推理速度，适用于移动端资源受限场景。

   quantized_model = torch.quantization.quantize_dynamic(
       traced_model, {torch.nn.Linear}, dtype=torch.qint8
   )
   quantized_model.save("model_quantized.pt")
   

2. 硬件适配优化

   • 针对Android设备，可利用PyTorch Mobile的XNNPACK库优化CPU推理。
   • 若目标设备为iOS，需确保模型兼容Core ML格式（需通过coremltools转换，见后续步骤）。

三、移动端部署

Android平台

1. 集成PyTorch Mobile库
   在build.gradle中添加依赖：

   implementation 'org.pytorch:pytorch_android:1.13.0'
   implementation 'org.pytorch:pytorch_android_torchvision:1.13.0'
   

2. 加载模型与推理

   // 加载模型
   Module module = Module.load("model.pt");
   
   // 准备输入（以图像为例）
   Tensor inputTensor = TensorImageUtils.bitmapToFloat32Tensor(
       bitmap, 
       TensorImageUtils.TORCHVISION_NORM_MEAN_RGB,
       TensorImageUtils.TORCHVISION_NORM_STD_RGB
   );
   
   // 执行推理
   Tensor outputTensor = module.forward(IValue.from(inputTensor)).toTensor();
   float[] results = outputTensor.getDataAsFloatArray();
   

iOS平台

1. 转换模型为Core ML格式
   通过coremltools将TorchScript模型转换为Core ML可识别的.mlmodel格式。

   import coremltools as ct
   
   example_input = torch.rand(1, 3, 224, 224)
   traced_model = torch.jit.trace(model, example_input)
   mlmodel = ct.convert(
       traced_model,
       inputs=[ct.TensorType(shape=example_input.shape)]
   )
   mlmodel.save("model.mlmodel")
   

2. Xcode集成与推理
   在Xcode项目中添加.mlmodel文件，通过MLModel类加载模型并执行推理。

四、注意事项

• 硬件加速：Android设备可启用GPU加速（需设备支持），iOS设备利用Metal框架优化。
• 性能测试：在真实设备上测试推理速度和内存占用，必要时调整模型结构或量化参数。
• 工具链依赖：确保Ubuntu环境安装libtorch、onnx等工具，可通过pip install torch torchvision获取。

参考资料：