TVM部署TensorFlow模型

本文介绍了如何使用TVM编译和部署TensorFlow模型,包括环境配置、模型导入、图像解码、Relay转换、编译执行和结果处理。在实践中遇到了TensorFlow版本问题,通过回退版本解决。完整教程和代码可在TVM官方文档中找到。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

本文是对TVM官方关于如何使用TVM编译TensorFlow模型文档的翻译整理,并记录了实现时遇到的小坑。

TVM部署TensorFlow模型

本文介绍如何使用TVM部署TensorFlow模型。
在开始之前,首先需要安装TensorFlow的Python包。

Python及TensorFlow环境

以下例程需要Python3.5以上的环境才能运行,请使用3.5以上Python的版本。
我在使用Python3.6.5与TensorFlow1.14.0运行例程时遇到如下报错:

Traceback (most recent call last):
  File "from_tensorflow.py", line 129, in <module>
    shape=shape_dict)
  File "/home/$USER/local/tvm/python/tvm/relay/frontend/tensorflow.py", line 2413, in from_tensorflow
    mod, params = g.from_tensorflow(graph, layout, shape, outputs)
  File "/home/$USER/local/tvm/python/tvm/relay/frontend/tensorflow.py", line 2058, in from_tensorflow
    op = self._convert_operator(node.op, inputs, attr, graph)
  File "/home/$USER/local/tvm/python/tvm/relay/frontend/tensorflow.py", line 2376, in _convert_operator
    sym = convert_map[op_name](inputs, attrs, self._params)
  File "/home/$USER/local/tvm/python/tvm/relay/frontend/tensorflow.py", line 562, in _impl
    extras={'method': "BILINEAR"})(inputs, attr)
  File "/home/$USER/local/tvm/python/tvm/relay/frontend/tensorflow.py", line 155, in __call__
    return _get_relay_op(op_name)(*inputs, **new_attrs)
TypeError: resize() got an unexpected keyword argument 'half_pixel_centers'

回退TensorFlow版本至1.12.0后报错消失。TVM社区关于本问题的探讨:https://discuss.tvm.ai/t/typeerror-when-running-the-from-tensorflow-example/3046

关于TensorFlow的安装请参考https://www.tensorflow.org/install
此处通过pip进行安装即可:

pip3 install TensorFlow==1.12.0
# 导入 tvm, relay  
import tvm  
from tvm import relay  

# 导入 os and numpy  
import numpy as np  
import os.path  

# 导入 Tensorflow imports  
import tensorflow as tf  

# Tensorflow 效用函数
import tvm.relay.testing.tf as tf_testing  

# 相关文件的在线地址(此处使用了dmlc在GitHub上的数据)  
repo_base = 'https://github.com/dmlc/web-data/raw/master/tensorflow/models/InceptionV1/'  

# 测试用图  
img_name = 'elephant-299.jpg'  
image_url = os.path.join(repo_base, img_name)  

教程

有关TensorFlow的各种模型的更多详细信息,请参阅 docs/frontend/tensorflow.md 。

model_name = 'classify_image_graph_def-with_shapes.pb'  
model_url = os.path.join(repo_base, model_name)  

# 图像标签  
map_proto = 'imagenet_2012_challenge_label_map_proto.pbtxt'  
map_proto_url = os.path.join(repo_base, map_proto)  

# 可读的图像标签  
label_map = 'imagenet_synset_to_human_label_map.txt'  
label_map_url = os.path.join(repo_base, label_map)  

# 目标设置  
# 如果使用cuda,可以使用以下推荐配置。  
#target = 'cuda'  
#target_host = 'llvm'  
#layout = "NCHW"  
#ctx = tvm.gpu(0)  
target = 'llvm'  
target_host = 'llvm'  
layout = None  
ctx = tvm.cpu(0)  

下载所需文件

下列程序将下载上面所列的所需文件

from tvm.contrib.download import download_testdata  

img_path = download_testdata(image_url, img_name, module='data')  
model_path = download_testdata(model_url, model_name, module=['tf', 'InceptionV1'])  
map_proto_path = download_testdata(map_proto_url, map_proto, module='data')  
label_path = download_testdata(label_map_url, label_map, module='data')  

导入模型

从protobuf文件创建TensorFlow图定义

with tf.gfile.FastGFile(model_path, 'rb') as f:  
    graph_def = tf.GraphDef()  
    graph_def.ParseFromString(f.read())  
    graph = tf.import_graph_def(graph_def, name='')  
    # 调用效用函数,将图定义导入默认图。  
    graph_def = tf_testing.ProcessGraphDefParam(graph_def)  
    # 向图增加shape。  
    with tf.Session() as sess:  
        graph_def = tf_testing.AddShapesToGraphDef(sess, 'softmax')  

图像解码

官方注解

TensorFlow前端导入不支持JpegDecode等处理操作,所以我们绕过JpegDecode(只返回源节点)。因此,我们需要向TVM提供已解码的帧。

from PIL import Image  
image = Image.open(img_path).resize((299, 299))  

x = np.array(image)  
</
### 如何使用 TVM 郰署 Whisper 模型 TVM 是一种开源机器学习编译器栈,支持多种硬件平台上的高效模型部署。以下是关于如何使用 TVM 部署 Whisper 模型的相关说明: #### 1. 安装 TVM 和依赖项 为了成功部署 Whisper 模型,首先需要安装 TVM 及其相关依赖库。可以参考官方文档完成环境搭建[^6]。 ```bash pip install tvm ``` 如果 Whisper 使用 PyTorch 或 TensorFlow 构建,则还需要确保这些框架已正确安装并兼容当前版本的 TVM。 #### 2. 导出模型结构与参数 Whisper 模型通常基于 Hugging Face Transformers 库实现。因此可以通过以下方式导出 ONNX 格式的模型文件以便后续导入至 TVM 中处理: ```python from transformers import AutoModelForSpeechSeq2Seq, AutoTokenizer import torch.onnx model_name = "openai/whisper-tiny" model = AutoModelForSpeechSeq2Seq.from_pretrained(model_name) tokenizer = AutoTokenizer.from_pretrained(model_name) dummy_input = tokenizer.encode("Hello world", return_tensors="pt") torch.onnx.export( model, dummy_input, f"{model_name}.onnx", input_names=["input_ids"], output_names=["logits"], dynamic_axes={"input_ids": {0: "batch_size"}, "logits": {0: "batch_size"}}, opset_version=13, ) ``` 此脚本会生成一个名为 `whisper-tiny.onnx` 的文件用于下一步加载到 TVM 上面去解析执行[^7]。 #### 3. 加载ONNX模型进入TVM IRModule 一旦获得了 `.onnx` 文件之后就可以借助于 TVM 提供的功能将其转换成内部表示形式(IR),从而进一步优化或者直接编译为目标设备上可运行的形式。 ```python import tvm.relay as relay from tvm.contrib.download import download_testdata onnx_file_url = "https://example.com/path/to/your/model.whisper.onnx" local_path = download_testdata(onnx_file_url,"model.whisper.onnx",module="onnx") mod, params = relay.frontend.from_onnx(local_path, shape=None, freeze_params=True) ``` 上述代码片段展示了怎样通过网络下载预训练好的 whisper onnx 模型,并且把它转交给 Relay 来构建计算图谱以及初始化权重参数表单[^8]。 #### 4. 编译目标硬件特定代码 最后一步就是针对具体的目标平台(CPU/GPU/FPGA etc.) 进行定制化的性能调整和最终产物制作过程如下所示 : ```python target = 'llvm' # or cuda/vulkan/opencl depending upon your target device type. with tvm.transform.PassContext(opt_level=3): lib = relay.build(mod,target,params=params) lib.export_library("/path/to/save/lib.so") ``` 这样就完成了整个流程从获取源码直到产出可供嵌入式系统使用的共享对象(shared object).so文件为止[^9]. --- ###
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值