PyTorch模型转TensorFlow Lite的Android部署全流程指南

pytorch model → onnx → tensorflow → tensorflow lite → android app
     ↑              ↑            ↑               ↑            ↑
  训练环境      中间格式     转换工具      优化部署      移动应用

import torch
import torch.nn as nn
from torchvision import models

# 创建或加载预训练模型
def create_model(num_classes=10):
    model = models.resnet18(pretrained=true)
    model.fc = nn.linear(model.fc.in_features, num_classes)
    return model

# 加载训练好的模型
model = create_model(num_classes=10)
model.load_state_dict(torch.load('best_model.pth', map_location='cpu'))
model.eval()

import torch.onnx

# 创建示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出 onnx 模型
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    export_params=true,        # 存储训练参数
    opset_version=14,          # onnx 算子集版本
    do_constant_folding=true,  # 执行常量折叠优化
    input_names=['input'],     # 输入名
    output_names=['output'],   # 输出名
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

print("onnx model exported successfully!")

pip install onnx-tf tensorflow

import onnx
from onnx_tf.backend import prepare
import tensorflow as tf

# 加载 onnx 模型
onnx_model = onnx.load("model.onnx")

# 转换为 tensorflow
tf_rep = prepare(onnx_model)
tf_rep.export_graph("saved_model")

print("tensorflow savedmodel created successfully!")

import numpy as np

# 测试输入
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# pytorch 推理
with torch.no_grad():
    pytorch_output = model(torch.from_numpy(test_input)).numpy()

# tensorflow 推理
tf_model = tf.saved_model.load("saved_model")
tf_output = tf_model(tf.constant(test_input.transpose(0, 2, 3, 1))).numpy()

# 验证数值一致性
np.testing.assert_allclose(pytorch_output, tf_output, rtol=1e-3)
print("conversion verified successfully!")

import tensorflow as tf

# 加载 savedmodel
converter = tf.lite.tfliteconverter.from_saved_model("saved_model")

# 转换为 tflite
tflite_model = converter.convert()

# 保存模型
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

print("basic tflite model created!")

# 启用所有优化
converter.optimizations = [tf.lite.optimize.default]

# 量化配置（显著减小模型大小）
converter.representative_dataset = representative_data_gen
converter.target_spec.supported_ops = [
    tf.lite.opsset.tflite_builtins_int8,
    tf.lite.opsset.select_tf_ops
]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8

# 转换
quantized_tflite_model = converter.convert()

# 保存量化模型
with open('model_quantized.tflite', 'wb') as f:
    f.write(quantized_tflite_model)

print("quantized tflite model created!")

def representative_data_gen():
    """生成代表性数据用于量化"""
    for _ in range(100):
        # 使用真实数据或随机数据
        data = np.random.rand(1, 224, 224, 3).astype(np.float32)
        yield [data]

android {
    compilesdk 34
    
    defaultconfig {
        applicationid "com.example.imagedemo"
        minsdk 24  // tensorflow lite requires api 24+
        targetsdk 34
        versioncode 1
        versionname "1.0"
    }
    
    compileoptions {
        sourcecompatibility javaversion.version_1_8
        targetcompatibility javaversion.version_1_8
    }
    
    // 启用 viewbinding
    buildfeatures {
        viewbinding true
    }
}

dependencies {
    implementation 'org.tensorflow:tensorflow-lite:2.15.0'
    implementation 'org.tensorflow:tensorflow-lite-support:0.4.4'
    implementation 'org.tensorflow:tensorflow-lite-metadata:0.4.4'
    
    // 可选：gpu 加速
    implementation 'org.tensorflow:tensorflow-lite-gpu:2.15.0'
    
    // 可选：nnapi 加速
    implementation 'org.tensorflow:tensorflow-lite-support:0.4.4'
}

package com.example.imagedemo;

import android.content.context;
import android.graphics.bitmap;
import android.util.log;
import org.tensorflow.lite.interpreter;
import org.tensorflow.lite.support.common.fileutil;
import org.tensorflow.lite.support.image.imageprocessor;
import org.tensorflow.lite.support.image.tensorimage;
import org.tensorflow.lite.support.image.ops.resizeop;
import org.tensorflow.lite.support.label.tensorlabel;
import org.tensorflow.lite.support.tensorbuffer.tensorbuffer;
import java.io.ioexception;
import java.nio.mappedbytebuffer;
import java.util.list;
import java.util.map;

public class imageclassifier {
    private static final string tag = "imageclassifier";
    private static final int input_image_size = 224;
    private static final float image_mean = 0.0f;
    private static final float image_std = 255.0f;
    
    private interpreter tflite;
    private list<string> labels;
    private tensorimage inputimagebuffer;
    private tensorbuffer outputprobabilitybuffer;
    private imageprocessor imageprocessor;
    
    public imageclassifier(context context) throws ioexception {
        // 加载模型
        mappedbytebuffer model = fileutil.loadmappedfile(context, "model_quantized.tflite");
        tflite = new interpreter(model);
        
        // 加载标签（可选）
        labels = fileutil.loadlabels(context, "labels.txt");
        
        // 初始化输入输出缓冲区
        inputimagebuffer = new tensorimage(android.graphics.bitmap.config.rgb_565);
        outputprobabilitybuffer = tensorbuffer.createfixedsize(new int[]{1, 10}, 
            datatype.float32);
        
        // 图像预处理器
        imageprocessor = new imageprocessor.builder()
            .add(new resizeop(input_image_size, input_image_size, resizeop.resizemethod.bilinear))
            .build();
    }
    
    public map<string, float> classify(bitmap bitmap) {
        // 预处理图像
        inputimagebuffer.load(bitmap);
        tensorimage processedimage = imageprocessor.process(inputimagebuffer);
        
        // 执行推理
        tflite.run(processedimage.getbuffer(), outputprobabilitybuffer.getbuffer().rewind());
        
        // 获取结果
        tensorlabel tensorlabel = new tensorlabel(labels, outputprobabilitybuffer);
        return tensorlabel.getmapwithfloatvalue();
    }
    
    public void close() {
        if (tflite != null) {
            tflite.close();
            tflite = null;
        }
    }
}

package com.example.imagedemo;

import android.manifest;
import android.content.intent;
import android.content.pm.packagemanager;
import android.graphics.bitmap;
import android.net.uri;
import android.os.bundle;
import android.provider.mediastore;
import android.util.log;
import android.widget.button;
import android.widget.imageview;
import android.widget.textview;
import androidx.annotation.nonnull;
import androidx.appcompat.app.appcompatactivity;
import androidx.core.app.activitycompat;
import androidx.core.content.contextcompat;
import java.io.ioexception;
import java.util.map;

public class mainactivity extends appcompatactivity {
    private static final string tag = "mainactivity";
    private static final int request_image = 1;
    private static final int request_permission = 2;
    
    private imageclassifier classifier;
    private imageview imageview;
    private textview resulttextview;
    
    @override
    protected void oncreate(bundle savedinstancestate) {
        super.oncreate(savedinstancestate);
        setcontentview(r.layout.activity_main);
        
        imageview = findviewbyid(r.id.imageview);
        resulttextview = findviewbyid(r.id.resulttextview);
        button selectimagebutton = findviewbyid(r.id.selectimagebutton);
        
        // 初始化分类器
        try {
            classifier = new imageclassifier(this);
        } catch (ioexception e) {
            log.e(tag, "failed to initialize classifier", e);
        }
        
        selectimagebutton.setonclicklistener(v -> selectimage());
    }
    
    private void selectimage() {
        if (contextcompat.checkselfpermission(this, manifest.permission.read_external_storage)
            != packagemanager.permission_granted) {
            activitycompat.requestpermissions(this,
                new string[]{manifest.permission.read_external_storage}, request_permission);
        } else {
            intent intent = new intent(intent.action_pick, mediastore.images.media.external_content_uri);
            startactivityforresult(intent, request_image);
        }
    }
    
    @override
    protected void onactivityresult(int requestcode, int resultcode, intent data) {
        super.onactivityresult(requestcode, resultcode, data);
        if (requestcode == request_image && resultcode == result_ok && data != null) {
            try {
                uri imageuri = data.getdata();
                bitmap bitmap = mediastore.images.media.getbitmap(getcontentresolver(), imageuri);
                imageview.setimagebitmap(bitmap);
                
                // 执行分类
                map<string, float> results = classifier.classify(bitmap);
                displayresults(results);
                
            } catch (ioexception e) {
                log.e(tag, "error processing image", e);
            }
        }
    }
    
    private void displayresults(map<string, float> results) {
        stringbuilder builder = new stringbuilder();
        for (map.entry<string, float> entry : results.entryset()) {
            builder.append(string.format("%s: %.2f%%\n", 
                entry.getkey(), entry.getvalue() * 100));
        }
        resulttextview.settext(builder.tostring());
    }
    
    @override
    protected void ondestroy() {
        super.ondestroy();
        if (classifier != null) {
            classifier.close();
        }
    }
}

<?xml version="1.0" encoding="utf-8"?>
<linearlayout xmlns:android="http://schemas.android.com/apk/res/android"
    android:layout_width="match_parent"
    android:layout_height="match_parent"
    android:orientation="vertical"
    android:padding="16dp">
    <imageview
        android:id="@+id/imageview"
        android:layout_width="match_parent"
        android:layout_height="300dp"
        android:scaletype="centercrop"
        android:background="#eeeeee" />
    <button
        android:id="@+id/selectimagebutton"
        android:layout_width="match_parent"
        android:layout_height="wrap_content"
        android:layout_margintop="16dp"
        android:text="select image" />
    <scrollview
        android:layout_width="match_parent"
        android:layout_height="0dp"
        android:layout_weight="1"
        android:layout_margintop="16dp">
        <textview
            android:id="@+id/resulttextview"
            android:layout_width="match_parent"
            android:layout_height="wrap_content"
            android:text="results will appear here"
            android:textsize="16sp" />
    </scrollview>
</linearlayout>

// 在 imageclassifier 中添加
private interpreter.options getgpuoptions() {
    interpreter.options options = new interpreter.options();
    gpudelegate gpudelegate = new gpudelegate();
    options.adddelegate(gpudelegate);
    return options;
}

// 使用 gpu 选项创建解释器
tflite = new interpreter(model, getgpuoptions());

// nnapi 选项
private interpreter.options getnnapioptions() {
    interpreter.options options = new interpreter.options();
    nnapidelegate nnapidelegate = new nnapidelegate();
    options.adddelegate(nnapidelegate);
    return options;
}

// 单例模式避免重复加载
public class classifiermanager {
    private static imageclassifier instance;
    
    public static synchronized imageclassifier getinstance(context context) {
        if (instance == null) {
            try {
                instance = new imageclassifier(context);
            } catch (ioexception e) {
                log.e("classifiermanager", "failed to create classifier", e);
            }
        }
        return instance;
    }
}

// 使用 asynctask 或 executorservice
private void classifyasync(bitmap bitmap) {
    executorservice executor = executors.newsinglethreadexecutor();
    handler handler = new handler(looper.getmainlooper());
    
    executor.execute(() -> {
        map<string, float> results = classifier.classify(bitmap);
        handler.post(() -> displayresults(results));
    });
}

# 使用 opset_version=14
torch.onnx.export(..., opset_version=14)

# 或者自定义操作替换
class custommodel(nn.module):
    def forward(self, x):
        # 避免使用不支持的操作
        return torch.clamp(x, 0, 1)  # 而不是 f.relu6

# 确保预处理一致
# pytorch: transforms.normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# tflite: 在 representative_data_gen 中使用相同归一化

// 确保正确的 abi 支持
android {
    defaultconfig {
        ndk {
            abifilters 'arm64-v8a', 'armeabi-v7a'
        }
    }
}

// 分离模型文件
android {
    bundle {
        language {
            enablesplit = false
        }
        density {
            enablesplit = false
        }
        abi {
            enablesplit = true  // 按 abi 分离
        }
    }
}

// 支持多图同时推理
public float[][] classifybatch(bitmap[] bitmaps) {
    int batchsize = bitmaps.length;
    tensorimage[] inputs = new tensorimage[batchsize];
    
    for (int i = 0; i < batchsize; i++) {
        inputs[i] = new tensorimage(bitmap.config.rgb_565);
        inputs[i].load(bitmaps[i]);
        inputs[i] = imageprocessor.process(inputs[i]);
    }
    
    // 批量推理
    object[] inputarray = arrays.stream(inputs)
        .map(tensorimage::getbuffer)
        .toarray(buffer[]::new);
    
    float[][] outputs = new float[batchsize][10];
    tflite.runformultipleinputsoutputs(inputarray, 
        new hashmap<integer, object>() {{
            put(0, outputs);
        }});
    
    return outputs;
}

// 在 assets 目录中包含模型元数据
// model_metadata.json
{
    "version": "1.2.0",
    "input_shape": [1, 224, 224, 3],
    "output_classes": 10,
    "preprocessing": {
        "mean": [0.485, 0.456, 0.406],
        "std": [0.229, 0.224, 0.225]
    }
}

// 支持多个模型文件
public class modelmanager {
    private static final string[] model_names = {
        "model_v1.tflite",
        "model_v2.tflite"
    };
    
    public imageclassifier getclassifier(context context, int version) {
        // 根据实验组选择模型
        return new imageclassifier(context, model_names[version]);
    }
}