Tensorflow(pb) 转 ONNX#
下面以 mobilenet_v2 float_v2_1.4_224 为例，展示 Tensorflow pb 模型转换为 ONNX 模型的过程:
import numpy as np
import tensorflow as tf
from tensorflow.core.framework.graph_pb2 import GraphDef
2024-01-08 20:09:39.954850: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-01-08 20:09:40.104429: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2024-01-08 20:09:40.104461: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
2024-01-08 20:09:40.139128: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-01-08 20:09:40.732640: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2024-01-08 20:09:40.732741: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2024-01-08 20:09:40.732753: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
import set_env # 加载 TVM
import tvm.relay.testing.tf as tf_testing
class MobilenetV2(tf.keras.Model):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
path_model = tf_testing.get_workload_official(
"https://storage.googleapis.com/mobilenet_v2/checkpoints/mobilenet_v2_1.4_224.tgz",
"mobilenet_v2_1.4_224_frozen.pb"
)
self.graph_def = self._read_graph_def(path_model)
self.output_names = ['output']
def _read_graph_def(self, frozen_path):
with open(frozen_path, 'rb') as f:
graph_def = GraphDef()
graph_def.ParseFromString(f.read())
return graph_def
@tf.function(input_signature=[tf.TensorSpec([1, 3, 224, 224],
tf.float32, name="input")])
def call(self, x):
x = tf.convert_to_tensor(x, tf.float32) # 确保输入是 tensor
x = tf.transpose(x, perm=(0, 2, 3, 1)) # NCHW -> NHWC
x = tf.graph_util.import_graph_def(
self.graph_def, input_map={'input:0': x},
return_elements=['MobilenetV2/Predictions/Reshape_1:0']
# return_elements=["MobilenetV2/Logits/AvgPool:0"]
# return_elements=["MobilenetV2/Logits/Conv2d_1c_1x1/BiasAdd:0"]
)[0]
return x
import tf2onnx
import onnx
input_signature = [tf.TensorSpec([1, 3, 224, 224], tf.float32, name="data")]
model = MobilenetV2()
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature)
onnx.save(onnx_model, ".temp/mobilenet_v2_tf.onnx")
2024-01-08 20:09:50.868403: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2024-01-08 20:09:50.868555: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory
2024-01-08 20:09:50.868644: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublasLt.so.11'; dlerror: libcublasLt.so.11: cannot open shared object file: No such file or directory
2024-01-08 20:09:50.911288: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusparse.so.11'; dlerror: libcusparse.so.11: cannot open shared object file: No such file or directory
2024-01-08 20:09:50.911367: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1934] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
2024-01-08 20:09:50.912367: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-01-08 20:09:53.929641: I tensorflow/core/grappler/devices.cc:66] Number of eligible GPUs (core count >= 8, compute capability >= 0.0): 2
2024-01-08 20:09:53.929931: I tensorflow/core/grappler/clusters/single_machine.cc:358] Starting new session
2024-01-08 20:09:54.247997: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1934] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
2024-01-08 20:09:55.192422: I tensorflow/core/grappler/devices.cc:66] Number of eligible GPUs (core count >= 8, compute capability >= 0.0): 2
2024-01-08 20:09:55.192620: I tensorflow/core/grappler/clusters/single_machine.cc:358] Starting new session
2024-01-08 20:09:55.193352: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1934] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
测试一致性:
from PIL import Image
image_size = 224
path = 'images/Giant_Panda_in_Beijing_Zoo_1.jpg' # 将要预测的图片路径
with Image.open(path) as im:
if im.mode != "RGB":
im.convert("RGB")
im = im.resize((224, 224))
image = np.asarray(im)
image = image/128 -1
images = np.expand_dims(image, 0)
images = images.transpose((0, 3, 1, 2))
model = MobilenetV2()
tf_output = model(images)
model.summary()
Model: "mobilenet_v2_1"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
=================================================================
Total params: 0
Trainable params: 0
Non-trainable params: 0
_________________________________________________________________
import set_env
import tvm
from tvm import relay
from tvm.relay.frontend import from_onnx
shape_dict = {"data": [1, 3, 224, 224]}
mod, params = from_onnx(
onnx_model,
shape_dict,
freeze_params=True
)
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, "llvm", params=params)
inputs_dict = {"data": images}
mlib_proxy = tvm.contrib.graph_executor.GraphModule(lib["default"](tvm.cpu()))
mlib_proxy.run(**inputs_dict)
np.testing.assert_allclose(
tf_output.numpy(),
mlib_proxy.get_output(0).numpy(),
rtol=1e-07, atol=1e-5
)
One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.